TORCH.JIT.UNUSED 's code can't run

[dugudongfangshuo]

📚 Documentation

document url is https://pytorch.org/docs/stable/generated/torch.jit.unused.html
in there code is

import torch
import torch.nn as nn

class MyModule(nn.Module):
    def __init__(self, use_memory_efficent):
        super(MyModule, self).__init__()
        self.use_memory_efficent = use_memory_efficent

    @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_efficent=False))
m.save("m.pt")

m = torch.jit.script(MyModule(use_memory_efficient=True))
# exception raised
m(torch.rand(100))

in the init the var is use_memory_efficent
but in def forward the var is use_memory_efficient
so this code can't run
correct is all use_memory_efficent
i hope you can fix it

cc @gmagogsfm

[eellison]

Good catch, thanks for the issue!