edits

curl_sac.py

@@ -482,6 +482,11 @@ class CurlSacAgent(object):
             self.critic.state_dict(), '%s/critic_%s.pt' % (model_dir, step)
         )
 
+    def save_curl(self, model_dir, step):
+        torch.save(
+            self.CURL.state_dict(), '%s/curl_%s.pt' % (model_dir, step)
+        )
+
     def load(self, model_dir, step):
         self.actor.load_state_dict(
             torch.load('%s/actor_%s.pt' % (model_dir, step))

train.py

@@ -230,7 +230,7 @@ def main():
             L.log('eval/episode', episode, step)
             evaluate(env, agent, video, args.num_eval_episodes, L, step,args)
             if args.save_model:
-                agent.save(model_dir, step)
+                agent.save_curl(model_dir, step)
             if args.save_buffer:
                 replay_buffer.save(buffer_dir)