PyTorch-lightning models running out of Memory after 1st epoch - deep-learning

Question

I saw a Kaggle kernel on PyTorch and run it with the same img_size, batch_size, etc. and created another PyTorch-lightning kernel with exact same values but my lightning model runs out of memory after about 1.5 epochs (each epoch contains 8750 steps) on the first fold whereas the native PyTorch model runs for whole 5 folds. Is there any way to improve the code or release memory? I could have tried to delete the models or do some garbage collection but if it doesn't complete even the first fold I can't delete the models and things.
def run_fold(fold):
df_train = train[train['fold'] != fold]
df_valid = train[train['fold'] == fold]
train_dataset = G2NetDataset(df_train, get_train_aug())
valid_dataset = G2NetDataset(df_valid, get_test_aug())
train_dl = DataLoader(train_dataset,
batch_size = config.batch_size,
num_workers = config.num_workers,
shuffle = True,
drop_last = True,
pin_memory = True)
valid_dl = DataLoader(valid_dataset,
batch_size = config.batch_size,
num_workers = config.num_workers,
shuffle = False,
drop_last = False,
pin_memory = True)
model = Classifier()
logger = pl.loggers.WandbLogger(project='G2Net', name=f'fold: {fold}')
trainer = pl.Trainer(gpus = 1,
max_epochs = config.epochs,
fast_dev_run = config.debug,
logger = logger,
log_every_n_steps=10)
trainer.fit(model, train_dl, valid_dl)
result = trainer.test(test_dataloaders = valid_dl)
wandb.run.finish()
return result
def main():
if config.train:
results = []
for fold in range(config.n_fold):
result = run_fold(fold)
results.append(result)
return results
results = main()

Answer 1

I cannot say much without looking at your model class, but couple possible issues that I encountered were metric and loss evaluation for logging.
For example, stuff like
pl.metrics.Accuracy(compute_on_step=False)
requires and explicit call of .compute()
def training_epoch_end(self, outputs):
loss = sum([out['loss'] for out in outputs])/len(outputs)
self.log_dict({'train_loss' : loss.detach(),
'train_accuracy' : self.train_metric.compute()})
at the epoch end.