commit from tesla

app.py

@@ -44,7 +44,7 @@ def train_model(model, tokenizer, data, device):
         num_train_epochs=3,
         per_device_train_batch_size=1,
         gradient_accumulation_steps=2,
-        # fp16=True,  # Enable mixed precision
+        fp16=True,  # Enable mixed precision
         warmup_steps=500,
         weight_decay=0.01,
         logging_dir='./logs',
@@ -61,11 +61,13 @@ def train_model(model, tokenizer, data, device):
     trainer.train()
 
      # Optionally clear cache if using GPU or MPS
+    print(torch.cuda.memory_summary(device=None, abbreviated=False))
+
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
     elif torch.has_mps:
         torch.mps.empty_cache()
-        
+
     # Perform any remaining steps such as logging, saving, etc.
     trainer.save_model()