Mirrors
/
waifu-diffusion
mirror of https://github.com/harubaru/waifu-diffusion.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

remove outputs from cudacallback

This commit is contained in:
harubaru 2022-09-07 03:51:04 -07:00 committed by GitHub
parent f7bb9b7bbc
commit d283e04b8b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
main.py
View File

@ -400,7 +400,7 @@ class CUDACallback(Callback):
torch.cuda.synchronize(trainer.root_gpu) torch.cuda.synchronize(trainer.root_gpu)
self.start_time = time.time() self.start_time = time.time()
def on_train_epoch_end(self, trainer, pl_module, outputs): def on_train_epoch_end(self, trainer, pl_module):
torch.cuda.synchronize(trainer.root_gpu) torch.cuda.synchronize(trainer.root_gpu)
max_memory = torch.cuda.max_memory_allocated(trainer.root_gpu) / 2 ** 20 max_memory = torch.cuda.max_memory_allocated(trainer.root_gpu) / 2 ** 20
epoch_time = time.time() - self.start_time epoch_time = time.time() - self.start_time