mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-09 11:32:07 +08:00
if return -1 from a hook that loop stopps
This commit is contained in:
@@ -327,7 +327,7 @@ class Trainer(TrainerIO):
|
||||
|
||||
# hook
|
||||
if self.__is_function_implemented('on_batch_start'):
|
||||
response = self.model.on_batch_start()
|
||||
response = self.model.on_batch_start(data_batch)
|
||||
if response == -1:
|
||||
return
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import torch
|
||||
|
||||
class ModelHooks(torch.nn.Module):
|
||||
def on_batch_start(self):
|
||||
def on_batch_start(self, data_batch):
|
||||
pass
|
||||
|
||||
def on_batch_end(self):
|
||||
|
||||
Reference in New Issue
Block a user