Simona/onnx (#161)

Creates a copy of the policy for exporting to avoid the interference of
the training and inference loops

emote/extra/onnx_exporter.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import io
 import logging
 import time
@@ -152,6 +153,11 @@ def process_pending_exports(self):
             item.process(self)
 
     def _trace(self):
+        with self.scopes.scope("policycopy"):
+            policy = copy.deepcopy(self.policy)
+
+        policy.train(False)
+
         with self.scopes.scope("trace"):
             args = []
 
@@ -162,15 +168,13 @@ def _trace(self):
 
                 args.append(arg)
 
-            self.policy.train(False)
-
             # NOTE: This might seem like a good use case for torch.jit.trace,
             # but it unfortunately leaks a full copy of the neural network.
             # See: https://github.com/pytorch/pytorch/issues/82532
 
             with io.BytesIO() as f:
                 torch.onnx.export(
-                    model=self.policy,
+                    model=policy,
                     args=tuple(args),
                     f=f,
                     input_names=list(map(lambda pair: pair[0], self.inputs)),
@@ -182,8 +186,6 @@ def _trace(self):
                     opset_version=13,
                 )
 
-                self.policy.train(True)
-
                 f.seek(0)
                 model_proto = onnx.load_model(f, onnx.ModelProto)