Merge branch 'master' of github.com:jackmcrider/zennit

- merge changes coming from github web interface

src/zennit/canonizers.py

@@ -414,22 +414,15 @@ def register(self, distance_module):
         weight = 2 * (self.distance.centroids[:, None, :] - self.distance.centroids[None, :, :])
         weight = weight[mask].reshape(n_clusters, n_clusters - 1, n_dims)
         norms = torch.norm(self.distance.centroids, dim=-1)
-        bias = (norms[None, :]**2 - norms[:, None]**2)[mask].reshape(n_clusters, n_clusters - 1)
-        setattr(self.parent_module, self.child_name,
-                torch.nn.Sequential(NeuralizedKMeans(weight, bias),
-                                    MinPool1d(n_clusters - 1),
-                                    torch.nn.Flatten()))
+        bias = (norms[None, :] ** 2 - norms[:, None] ** 2)[mask].reshape(n_clusters, n_clusters - 1)
+        self.parent_module.add_module(
+            self.child_name,
+            torch.nn.Sequential(NeuralizedKMeans(weight, bias), MinPool1d(n_clusters - 1), torch.nn.Flatten())
+        )
 
     def remove(self):
         """Revert the changes introduced by this canonizer."""
         setattr(self.parent_module, self.child_name, self.distance)
 
     def copy(self):
-        '''Copy this Canonizer.
-
-        Returns
-        -------
-        :py:obj:`Canonizer`
-            A copy of this Canonizer.
-        '''
         return KMeansCanonizer()

src/zennit/layer.py

@@ -114,8 +114,7 @@ def forward(self, x):
         :py:obj:`torch.Tensor`
             shape (N, K, K-1) tensor of k-means discriminants
         '''
-        x = torch.einsum('nd,kjd->nkj', x, self.weight) + self.bias
-        return x
+        return torch.einsum('nd,kjd->nkj', x, self.weight) + self.bias
 
 
 class MinPool2d(torch.nn.MaxPool2d):