add more informative error message

lucidrains · Jul 15, 2021 · caa65d7 · caa65d7
1 parent 2aa84ee
commit caa65d7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/byol_pytorch/byol_pytorch.py b/byol_pytorch/byol_pytorch.py
@@ -231,6 +231,8 @@ def forward(
         return_embedding = False,
         return_projection = True
     ):
+        assert not (self.training and x.shape[0] == 1), 'you must have greater than 1 sample when training, due to the batchnorm in the projection layer'
+
         if return_embedding:
             return self.online_encoder(x, return_projection = return_projection)
 

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