torch.nn.Module Memory Address Changes After Importing torchdrug

[benyaminjami]

Hello,

I've encountered an unexpected behavior when using torchdrug alongside PyTorch. Specifically, after importing torchdrug, the memory address of torch.nn.Module changes, which leads to issues when checking instances of nn.Module.

Here's a simple reproduction of the issue:

import torch.nn as nn

# Print the memory address of nn.Module before importing torchdrug
print(id(nn.Module))  # Outputs: 104525552

import torchdrug

# Print the memory address of nn.Module after importing torchdrug
print(id(nn.Module))  # Outputs: 172015376

As observed, the memory address of nn.Module changes after importing torchdrug, which suggests that the torch.nn.Module class is being shadowed or replaced in some way by torchdrug.

This behavior leads to unexpected issues, such as the isinstance check failing for classes that inherit from nn.Module.

Questions:

Is this behavior intentional? If so, what's the rationale behind it? Are there any recommended workarounds or solutions to address this issue? Is this issue known, and are there any plans to address it in future releases? Thank you for your assistance and looking forward to your feedback.