fix issue with hidden representation coming from last layer

lucidrains · Jun 17, 2020 · 85997ee · 85997ee
1 parent 4d27acd
commit 85997ee
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/byol_pytorch/byol_pytorch.py b/byol_pytorch/byol_pytorch.py
@@ -122,17 +122,20 @@ def _get_projector(self, hidden):
         projector = MLP(dim, self.projection_size, self.projection_hidden_size)
         return projector.to(hidden)
 
-    def forward(self, x):
+    def get_representation(self, x):
         if self.layer == -1:
             return self.net(x)
 
         _ = self.net(x)
         hidden = self.hidden
         self.hidden = None
         assert hidden is not None, f'hidden layer {self.layer} never emitted an output'
+        return hidden
 
-        projector = self._get_projector(hidden)
-        projection = projector(hidden)
+    def forward(self, x):
+        representation = self.get_representation(x)
+        projector = self._get_projector(representation)
+        projection = projector(representation)
         return projection
 
 # main class

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'byol-pytorch',
   packages = find_packages(),
-  version = '0.0.3',
+  version = '0.0.5',
   license='MIT',
   description = 'Self-supervised contrastive learning made simple',
   author = 'Phil Wang',