From caa65d7c19a80a3611f4049c638ecd8401f8f3c4 Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Thu, 15 Jul 2021 10:58:42 -0700
Subject: [PATCH] add more informative error message

---
 byol_pytorch/byol_pytorch.py | 2 ++
 setup.py                     | 2 +-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/byol_pytorch/byol_pytorch.py b/byol_pytorch/byol_pytorch.py
index 171f478e9..055e4ca6f 100644
--- 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
index ba2208173..cbbeb3d8b 100644
--- 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',