Torchscript

foo_scripted = torch.jit.script(foo)

@torch.jit.script

def foo2(x, y):

if x.max() > y.max():

r = x

else: r = y

return r

# using a decorator

• TorchScript excepts static typing, dynamic changes of variables are not allowed

• If we want to script non-tensor objects, we have to include a type hint

my_module_scripted = torch.jit.script(MyModule(5))

# this will automatically script the forward() method

# and all methods invoked by it

foo_traced = torch.jit.trace(foo, example_inputs=(tensor1, tensor2))

• Tracing does not take care of control flow, and therefore only evaluates the code which was executed with the initial input

foo_traced.save("10_torch_script_module.pt")

loaded_script_module = torch.jit.load("10_torch_script_module.pt")

foo_scripted.graph