Fixing warning triggered by torch.cuda.reset_max_memory_allocated()…

… usage
pyg-team · Jan 10, 2025 · 78e15d0 · 78e15d0
1 parent ef02854
commit 78e15d0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/llm/molecule_gpt.py b/examples/llm/molecule_gpt.py
@@ -167,7 +167,7 @@ def adjust_learning_rate(param_group, LR, epoch):
                 f'moleculegpt_epoch{best_epoch}_val_loss{best_val_loss:4f}_ckpt.pt'  # noqa: E501
             )
     torch.cuda.empty_cache()
-    torch.cuda.reset_max_memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
 
     print(f"Total Training Time: {time.time() - start_time:2f}s")
     # Test