wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-11 12:31:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

made checkpoint callback optional

Browse Source
This commit is contained in:
William Falcon
2019-07-15 13:17:38 -04:00
parent ab00514ef6
commit e57f461323
1 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/models/trainer.py
+6 -4
View File
@@ -31,7 +31,8 @@ class Trainer(TrainerIO):
def __init__(self,
experiment,
checkpoint_callback, early_stop_callback,
early_stop_callback,
checkpoint_callback=None,
gradient_clip=0,
cluster=None,
process_position=0,
@@ -68,7 +69,7 @@ class Trainer(TrainerIO):
self.process_position = process_position
self.current_gpu_name = current_gpu_name
self.checkpoint_callback = checkpoint_callback
self.checkpoint_callback.save_function = self.save_checkpoint
self.checkpoint_callback.save_function = self.save_checkpoint if self.checkpoint_callback is not None else None
self.early_stop = early_stop_callback
self.model = None
self.max_nb_epochs = max_nb_epochs
@@ -720,5 +721,6 @@ class Trainer(TrainerIO):
# model checkpointing
if self.proc_rank == 0:
print('save callback...')
self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)
if self.checkpoint_callback:
print('save callback...')
self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)