fix projected coordinates dimensions

lucidrains · Dec 19, 2020 · 9ccf7e6 · 9ccf7e6
1 parent 1727fec
commit 9ccf7e6
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/pixel_level_contrastive_learning/pixel_level_contrastive_learning.py b/pixel_level_contrastive_learning/pixel_level_contrastive_learning.py
@@ -386,8 +386,10 @@ def forward(self, x):
         pdist = nn.PairwiseDistance(p = 2)
 
         num_pixels = proj_coors_one.shape[0]
-        proj_coors_one_expanded = proj_coors_one[None, :].expand(num_pixels, num_pixels, -1).reshape(num_pixels * num_pixels, 2)
-        proj_coors_two_expanded = proj_coors_two[:, None].expand(num_pixels, num_pixels, -1).reshape(num_pixels * num_pixels, 2)
+
+        proj_coors_one_expanded = proj_coors_one[:, None].expand(num_pixels, num_pixels, -1).reshape(num_pixels * num_pixels, 2)
+        proj_coors_two_expanded = proj_coors_two[None, :].expand(num_pixels, num_pixels, -1).reshape(num_pixels * num_pixels, 2)
+
         distance_matrix = pdist(proj_coors_one_expanded, proj_coors_two_expanded)
         distance_matrix = distance_matrix.reshape(num_pixels, num_pixels)
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'pixel-level-contrastive-learning',
   packages = find_packages(),
-  version = '0.0.10',
+  version = '0.0.11',
   license='MIT',
   description = 'Pixel-Level Contrastive Learning',
   author = 'Phil Wang',