[Bug] TorchScript jit trace AttributeError

[andrenguyen]

🐛 Bug

Hi! I tried to follow the instructions in the Converting Variational Models to TorchScript notebook with a SVDKL model I trained, but I am running into an Attribute Error.

To reproduce

wrapped_svdkl = MeanStdSVDKLWrapper(model, likelihood)
with torch.no_grad(), gpytorch.settings.trace_mode():
    pred = wrapped_svdkl(inputs)  
    traced_svdkl = torch.jit.trace(wrapped_svdkl, inputs)

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-53-8bb7201b88f0> in <module>
      2 with torch.no_grad(), gpytorch.settings.trace_mode():
      3     pred = wrapped_svdkl(inputs)
----> 4     traced_svdkl = torch.jit.trace(wrapped_svdkl, inputs)

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in trace(func, example_inputs, optimize, check_trace, check_inputs, check_tolerance, _force_outplace, _module_class, _compilation_unit)
    873         return trace_module(func, {'forward': example_inputs}, None,
    874                             check_trace, wrap_check_inputs(check_inputs),
--> 875                             check_tolerance, _force_outplace, _module_class)
    876 
    877     if (hasattr(func, '__self__') and isinstance(func.__self__, torch.nn.Module) and

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in trace_module(mod, inputs, optimize, check_trace, check_inputs, check_tolerance, _force_outplace, _module_class, _compilation_unit)
   1019         register_submods(mod, '__module')
   1020 
-> 1021         module = make_module(mod, _module_class, _compilation_unit)
   1022 
   1023         for method_name, example_inputs in inputs.items():

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in make_module(mod, _module_class, _compilation_unit)
    718         if _module_class is None:
    719             _module_class = TopLevelTracedModule
--> 720         return _module_class(mod, _compilation_unit=_compilation_unit)
    721 
    722 def wrap_check_inputs(check_inputs):

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in __init__(self, orig, id_set, _compilation_unit)
   1882 
   1883         for name, submodule in orig._modules.items():
-> 1884             tmp_module._modules[name] = make_module(submodule, TracedModule, _compilation_unit=None)
   1885 
   1886         script_module = torch.jit._recursive.create_script_module(tmp_module, lambda module: (), share_types=False)

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in make_module(mod, _module_class, _compilation_unit)
    718         if _module_class is None:
    719             _module_class = TopLevelTracedModule
--> 720         return _module_class(mod, _compilation_unit=_compilation_unit)
    721 
    722 def wrap_check_inputs(check_inputs):

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in __init__(self, orig, id_set, _compilation_unit)
   1882 
   1883         for name, submodule in orig._modules.items():
-> 1884             tmp_module._modules[name] = make_module(submodule, TracedModule, _compilation_unit=None)
   1885 
   1886         script_module = torch.jit._recursive.create_script_module(tmp_module, lambda module: (), share_types=False)

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in make_module(mod, _module_class, _compilation_unit)
    718         if _module_class is None:
    719             _module_class = TopLevelTracedModule
--> 720         return _module_class(mod, _compilation_unit=_compilation_unit)
    721 
    722 def wrap_check_inputs(check_inputs):

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in __init__(self, orig, id_set, _compilation_unit)
   1882 
   1883         for name, submodule in orig._modules.items():
-> 1884             tmp_module._modules[name] = make_module(submodule, TracedModule, _compilation_unit=None)
   1885 
   1886         script_module = torch.jit._recursive.create_script_module(tmp_module, lambda module: (), share_types=False)

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in make_module(mod, _module_class, _compilation_unit)
    718         if _module_class is None:
    719             _module_class = TopLevelTracedModule
--> 720         return _module_class(mod, _compilation_unit=_compilation_unit)
    721 
    722 def wrap_check_inputs(check_inputs):

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in __init__(self, orig, id_set, _compilation_unit)
   1882 
   1883         for name, submodule in orig._modules.items():
-> 1884             tmp_module._modules[name] = make_module(submodule, TracedModule, _compilation_unit=None)
   1885 
   1886         script_module = torch.jit._recursive.create_script_module(tmp_module, lambda module: (), share_types=False)

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in make_module(mod, _module_class, _compilation_unit)
    718         if _module_class is None:
    719             _module_class = TopLevelTracedModule
--> 720         return _module_class(mod, _compilation_unit=_compilation_unit)
    721 
    722 def wrap_check_inputs(check_inputs):

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in __init__(self, orig, id_set, _compilation_unit)
   1882 
   1883         for name, submodule in orig._modules.items():
-> 1884             tmp_module._modules[name] = make_module(submodule, TracedModule, _compilation_unit=None)
   1885 
   1886         script_module = torch.jit._recursive.create_script_module(tmp_module, lambda module: (), share_types=False)

~/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py in make_module(mod, _module_class, _compilation_unit)
    701     if isinstance(mod, ScriptModule):
    702         return mod
--> 703     elif torch._jit_internal.module_has_exports(mod):
    704         def make_stubs_from_exported_methods(mod):
    705             exported = []

~/anaconda3/lib/python3.7/site-packages/torch/_jit_internal.py in module_has_exports(mod)
    436 def module_has_exports(mod):
    437     for name in dir(mod):
--> 438         item = getattr(mod, name)
    439         if callable(item):
    440             if get_torchscript_modifier(item) is FunctionModifiers.EXPORT:

~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py in __getattr__(self, name)
    592                 return modules[name]
    593         raise AttributeError("'{}' object has no attribute '{}'".format(
--> 594             type(self).__name__, name))
    595 
    596     def __setattr__(self, name, value):

AttributeError: 'SmoothedBoxPrior' object has no attribute 'islazy'

Expected Behavior

A traced model was expected.

System information

• GPyTorch Version 1.1.1
• PyTorch Version 1.5.0
• CentOS 7

Additional context

Thank you!

[gpleiss]

@jacobrgardner