commit from tesla

app.py

@@ -60,6 +60,12 @@ def train_model(model, tokenizer, data, device):
     
     trainer.train()
 
+     # Optionally clear cache if using GPU or MPS
+    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()
 
@@ -70,8 +76,8 @@ def main(api_name, base_url):
     # Load the configuration for a specific model
     config = AutoConfig.from_pretrained('google/codegemma-2b')
     # Update the activation function
-    config.hidden_act = ''  # Set to use approximate GeLU gelu_pytorch_tanh
-    config.hidden_activation = 'gelu'  # Set to use GeLU
+    # config.hidden_act = ''  # Set to use approximate GeLU gelu_pytorch_tanh
+    config.hidden_activation = 'gelu_pytorch_tanh'  # Set to use GeLU
 
     model = AutoModelForCausalLM.from_pretrained('google/codegemma-2b', is_decoder=True)
     #model = BertLMHeadModel.from_pretrained('google/codegemma-2b', is_decoder=True)
@@ -80,7 +86,6 @@ def main(api_name, base_url):
     model.to(device)  # Move model to the appropriate device
 
     train_model(model, tokenizer, data, device)
- 
     
     model.save_pretrained("./fine_tuned_model")
     tokenizer.save_pretrained("./fine_tuned_model")