fixed byol

Lightning-Universe · Aug 22, 2020 · efa1db0 · efa1db0
1 parent c98bb5b
commit efa1db0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pl_bolts/models/self_supervised/byol/byol_module.py b/pl_bolts/models/self_supervised/byol/byol_module.py
@@ -100,7 +100,7 @@ def __init__(self,
 
     def on_train_batch_end(self, batch: Any, batch_idx: int, dataloader_idx: int) -> None:
         # Add callback for user automatically since it's key to BYOL weight update
-        self.weight_callback.on_batch_end(self.trainer, self)
+        self.weight_callback.on_train_batch_end(self.trainer, self, batch, batch_idx, dataloader_idx)
 
     def forward(self, x):
         y, _, _ = self.online_network(x)