Save models

Code/train.py

@@ -190,7 +190,7 @@ def train(model, name, imdb, seed,gradient_accumulation_steps,mixup, threshold,
 			#progress_bar.update(1)
 		#print("one epoch done")
 	
-	torch.save(model.state_dict(), "bert_baseline.pt")
+	torch.save(model, "./saved_models/bert_baseline.pt")
 
 	#print(model_name)
 	evaluation_test = evaluation.evaluate_model(model, name,  test_dataset, learning_rate, test_batch_size, imdb)