From 343af70a8da22bb8934d1242a54f91661007d113 Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Wed, 13 May 2020 09:21:13 -0600
Subject: [PATCH] Add code for compiling model to torchscript

I want to be able to export it to other formats too in the future.
---
 codes/distill_torchscript.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)
 create mode 100644 codes/distill_torchscript.py

diff --git a/codes/distill_torchscript.py b/codes/distill_torchscript.py
new file mode 100644
index 00000000..07f2ecd8
--- /dev/null
+++ b/codes/distill_torchscript.py
@@ -0,0 +1,17 @@
+import argparse
+import options.options as option
+from models.networks import define_G
+import torch
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-opt', type=str, help='Path to options YMAL file.', default='options/test/test_ESRGAN_adrianna_full.yml')
+    opt = option.parse(parser.parse_args().opt, is_train=False)
+    opt = option.dict_to_nonedict(opt)
+    netG = define_G(opt)
+
+    print("Tracing generator network..")
+    dummyInput = torch.rand(1, 3, 8, 8)
+    traced_netG = torch.jit.trace(netG, dummyInput)
+    traced_netG.save('../results/traced_generator.zip')
+    print(traced_netG)
\ No newline at end of file