Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

debug

This commit is contained in:
Adrian Wälchli 2021-07-01 20:34:39 +02:00
parent 856cd6623b
commit a6e6101946
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
pytorch_lightning/loops/epoch/training_epoch_loop.py
View File

@ -110,11 +110,13 @@ class TrainingEpochLoop(Loop):
# TRAINING_STEP + TRAINING_STEP_END
# ------------------------------------
with self.trainer.profiler.profile("training_batch_to_device"):
print("before run", self.iteration_count, torch.cuda.memory_allocated())
batch = self.trainer.accelerator.batch_to_device(batch, dataloader_idx=self._dataloader_idx)
with self.trainer.profiler.profile("run_training_batch"):
batch_output = self.batch_loop.run(batch, self.iteration_count, self._dataloader_idx)
self.batches_seen += 1
print("after run", self.iteration_count, torch.cuda.memory_allocated())
# when returning -1 from train_step, we end epoch early
if batch_output.signal == -1:
@ -150,6 +152,7 @@ class TrainingEpochLoop(Loop):
Raises:
StopIteration: if :attr:`done` evaluates to ``True`` to finish this epoch
"""
print("advance end", self.iteration_count, torch.cuda.memory_allocated())
# -----------------------------------------
# VALIDATE IF NEEDED + CHECKPOINT CALLBACK
# -----------------------------------------