torch.jit.unused
-
torch.jit.unused(fn)
[source] -
This decorator indicates to the compiler that a function or method should be ignored and replaced with the raising of an exception. This allows you to leave code in your model that is not yet TorchScript compatible and still export your model.
Example (using
@torch.jit.unused
on a method):import torch import torch.nn as nn class MyModule(nn.Module): def __init__(self, use_memory_efficient): super(MyModule, self).__init__() self.use_memory_efficient = use_memory_efficient @torch.jit.unused def memory_efficient(self, x): import pdb pdb.set_trace() return x + 10 def forward(self, x): # Use not-yet-scriptable memory efficient mode if self.use_memory_efficient: return self.memory_efficient(x) else: return x + 10 m = torch.jit.script(MyModule(use_memory_efficient=False)) m.save("m.pt") m = torch.jit.script(MyModule(use_memory_efficient=True)) # exception raised m(torch.rand(100))
© 2019 Torch Contributors
Licensed under the 3-clause BSD License.
https://pytorch.org/docs/1.8.0/generated/torch.jit.unused.html