Pytorch export trace/script

[theFilipko]

Hello. Has anyone tried to save the model with partial convolutions as pytorch script? I got this error when calling pytorch.jit.script:

Traceback (most recent call last):
  File "export.py", line 47, in <module>
    net_script = torch.jit.script(model)
  File "/home/user/.local/lib/python3.6/site-packages/torch/jit/_script.py", line 943, in script
    obj, torch.jit._recursive.infer_methods_to_compile
  File "/home/user/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 391, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/home/user/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 448, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/home/user/.local/lib/python3.6/site-packages/torch/jit/_script.py", line 391, in _construct
    init_fn(script_module)
  File "/home/user/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 428, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/home/user/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 448, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/home/user/.local/lib/python3.6/site-packages/torch/jit/_script.py", line 391, in _construct
    init_fn(script_module)
  File "/home/user/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 428, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/home/user/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 452, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/home/user/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 335, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError: 
cannot statically infer the expected size of a list in this context:
  File "/home/user/Projects/gitlab/2d22d-poc/pconv.py", line 56
    def forward(self, input, mask_in=None):
        assert len(input.shape) == 4, "Input shape for partial convolution must have 4 dimensions"
        if mask_in is not None or self.last_size != tuple(input.shape):
                                                    ~~~~~~~~~~~~~~~~~ <--- HERE
            self.last_size = tuple(input.shape)