Fix direction scaling

marcellodebernardi · Aug 30, 2019 · 8d34610 · 8d34610
1 parent 7c92488
commit 8d34610
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/loss_landscapes/main.py b/loss_landscapes/main.py
@@ -143,7 +143,7 @@ def random_line(model_start: typing.Union[torch.nn.Module, ModelWrapper], metric
     else:
         raise AttributeError('Unsupported normalization argument. Supported values are model, layer, and filter')
 
-    direction.mul_(((start_point.model_norm() / distance) / steps) / direction.model_norm())
+    direction.mul_(((start_point.model_norm() * distance) / steps) / direction.model_norm())
 
     data_values = []
     for i in range(steps):
@@ -290,8 +290,8 @@ def random_plane(model: typing.Union[torch.nn.Module, ModelWrapper], metric: Met
         raise AttributeError('Unsupported normalization argument. Supported values are model, layer, and filter')
 
     # scale to match steps and total distance
-    dir_one.mul_(((start_point.model_norm() / distance) / steps) / dir_one.model_norm())
-    dir_two.mul_(((start_point.model_norm() / distance) / steps) / dir_two.model_norm())
+    dir_one.mul_(((start_point.model_norm() * distance) / steps) / dir_one.model_norm())
+    dir_two.mul_(((start_point.model_norm() * distance) / steps) / dir_two.model_norm())
     # Move start point so that original start params will be in the center of the plot
     dir_one.mul_(steps / 2)
     dir_two.mul_(steps / 2)

diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
 
 setup(
     name='loss_landscapes',
-    version='3.0.6',
+    version='3.0.7',
     packages=find_packages(exclude='tests'),
     url='https://github.com/marcellodebernardi/loss-landscapes',
     license='MIT',