save when training complete

.gitignore

@@ -19,16 +19,24 @@
 *.rar
 *.tar
 *.zip
-*.csv
-*.gpickle
-*.graphml
-*.dat
 
 # Logs and databases #
 ######################
 *.log
 *.sql
 *.sqlite
+*.csv
+*.gpickle
+*.graphml
+
+# Models #
+######################
+*.pt
+
+# Temporary Files #
+######################
+*.pyc
+*.dat
 
 # OS generated files #
 ######################

main.py

@@ -12,6 +12,7 @@
 import sys
 import warnings
 import time
+import datetime
 from pprint import pprint
 
 import matplotlib
@@ -137,7 +138,7 @@ def train(args, env, memory, policy_net, target_net, optimizer):
             reward = torch.tensor([reward], device=device)
 
             # Store the transition in memory
-            
+
             memory.push(state, action, next_state, reward, goal_state_embedding.detach())
 
             # Move to the next state
@@ -155,8 +156,19 @@ def train(args, env, memory, policy_net, target_net, optimizer):
             target_net.load_state_dict(policy_net.state_dict())
 
     print('Training Complete')
-    env.render()
+    # env.render()
+
+    save_dict = {'state_dict': target_net.state_dict(), 'args': args}
+    dest_path = f"models/{args.wiki_year}_fixednode-{args.has_fixed_dest_node}_{datetime.datetime.now().strftime('%Y_%m_%d-%I:%M:%S_%p')}.pt"
+    torch.save(save_dict, dest_path)
+    print('Model saved to location ', dest_path)
+
     env.close()
+    # torch.save(model.state_dict(), filepath)
+
+    # #Later to restore:
+    # model.load_state_dict(torch.load(filepath))
+    # model.eval()
 
 def plot_durations(episode_durations):
     plt.figure(2)