Correcting default behavior

I am thinking we expect to only get the feature embeddings if we write: model(image, return_embedding = True). Maybe not? Just wondering :)

Author: gauenk

byol_pytorch/byol_pytorch.py

@@ -227,7 +227,7 @@ def update_moving_average(self):
 
     def forward(self, x, return_embedding = False):
         if return_embedding:
-            return self.online_encoder(x)
+            return self.online_encoder(x, True)
 
         image_one, image_two = self.augment1(x), self.augment2(x)