fixes a bug when adding 2 PVector objects when a layer has no bias

tfjgeorge · May 9, 2020 · 1f648ef · 1f648ef
1 parent 83fe96d
commit 1f648ef
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/nngeometry/object/pspace.py b/nngeometry/object/pspace.py
@@ -411,7 +411,7 @@ def vTMv(self, vector):
         for layer_id, layer in self.generator.layer_collection.layers.items():
             v = vector_dict[layer_id][0].view(vector_dict[layer_id][0].size(0),
                                               -1)
-            if len(vector_dict[layer_id]) > 1:
+            if layer.bias is not None:
                 v = torch.cat([v, vector_dict[layer_id][1].unsqueeze(1)],
                               dim=1)
             a, g = self.data[layer_id]

diff --git a/nngeometry/object/vector.py b/nngeometry/object/vector.py
@@ -50,7 +50,7 @@ def from_model(model):
             if layer.bias is not None:
                 dict_repr[layer_id] = (mod.weight, mod.bias)
             else:
-                dict_repr[layer_id] = (mod.weight)
+                dict_repr[layer_id] = (mod.weight,)
         return PVector(layer_collection, dict_repr=dict_repr)
 
     @staticmethod
@@ -140,7 +140,7 @@ def __rmul__(self, x):
                     v_dict[l_id] = (x * self.dict_repr[l_id][0],
                                     x * self.dict_repr[l_id][1])
                 else:
-                    v_dict[l_id] = (x * self.dict_repr[l_id][0])
+                    v_dict[l_id] = (x * self.dict_repr[l_id][0],)
             return PVector(self.layer_collection, dict_repr=v_dict)
         else:
             return PVector(self.layer_collection,
@@ -157,7 +157,7 @@ def __add__(self, other):
                                     other.dict_repr[l_id][1])
                 else:
                     v_dict[l_id] = (self.dict_repr[l_id][0] +
-                                    other.dict_repr[l_id][0])
+                                    other.dict_repr[l_id][0],)
             return PVector(self.layer_collection, dict_repr=v_dict)
         elif self.vector_repr is not None and other.vector_repr is not None:
             return PVector(self.layer_collection,
@@ -178,7 +178,7 @@ def __sub__(self, other):
                                     other.dict_repr[l_id][1])
                 else:
                     v_dict[l_id] = (self.dict_repr[l_id][0] -
-                                    other.dict_repr[l_id][0])
+                                    other.dict_repr[l_id][0],)
             return PVector(self.layer_collection, dict_repr=v_dict)
         elif self.vector_repr is not None and other.vector_repr is not None:
             return PVector(self.layer_collection,