Merge pull request #6 from phillies/fix/init_bits

jacobkimmel · web-flow · commit 60c98d30feb8 · 2020-06-03T09:52:16.000-07:00
fix bits parameter in init function
diff --git a/pytorch_modelsize.py b/pytorch_modelsize.py
@@ -12,7 +12,7 @@ def __init__(self, model, input_size=(1,1,32,32), bits=32):
         '''
         self.model = model
         self.input_size = input_size
-        self.bits = 32
+        self.bits = bits
         return
 
     def get_parameter_sizes(self):
