Conflict with torch due to overwritten modules

[jannisborn]

I'm interested to understand why it is necessary to overwrite the default nn.Module of torch in patch.py:

torchdrug/torchdrug/patch.py

Line 125 in eeee191

nn.Module = PatchedModule

This seems to be a quite invasive thing since it alters the behavior of any torch.nn module after torchdrug has been imported.

For example, your implementation of register_buffer in patch.py lacks the keyword argument persistent which is present in native torch: https://github.com/pytorch/pytorch/blob/989b24855efe0a8287954040c89d679625dcabe1/torch/nn/modules/module.py#L277

I would greatly appreciate if you could please let me know how I can fall back to the native torch behavior after having imported torchdrug somewhere above in my code.

[KiddoZhu]

Hi! The patch is intended to make nn.Module and nn.DistributedDataParallel to accept data.Graph (and its derived classes) in the same way as torch.Tensor. For example, we may want to register a graph as a buffer of nn.Module, or pass graphs as arguments to some nn.Module.forward in distributed data parallel.

The code was originally developed based on PyTorch 1.4.0 -- and that's why persistent is missing there. I am aware that such a monkey patch is really ugly and not good for maintainence. Any idea we can improve that?

[KiddoZhu]

The original nn.Module is reassigned to nn._Module in the patch. To fall back to native PyTorch, you may use

torch.nn.Module = torch.nn._Module

But I guess this will make many TorchDrug classes complain.

[KiddoZhu]

I noticed that persistent is always an argument in register_buffer for PyTorch>=1.6.0. So I updated the corresponding patch function in b411837. Feel free to reopen this issue if you have any better solution for such patches.

[jannisborn]

Cool, thanks a lot for looking into it!