Unable to reload class inherited from nn.Module with importlib
I am trying to reload a class using importlib, however I am facing an error stating it is not a module, this is an jupyter notebook.
Class code
import torch.nn as nn
import torch.nn.functional as F
import torch

class FeedForwardNeuralNetwork(nn.Module):
    def __init__(self, input_size, layers_data, random_seed=42):
        super(FeedForwardNeuralNetwork, self).__init__()
        torch.manual_seed(random_seed)
        # So that number of dense layers are configurable.
        self.layers = nn.ModuleList()
        for size, activation in layers_data:
            self.layers.append(nn.Linear(input_size, size, bias=False))
            torch.nn.init.xavier_uniform(self.layers[-1].weight)
            self.layers.append(activation)
            input_size = size 
    def forward(self, x):
        x = torch.flatten(x, start_dim=1)
        for layer in self.layers:
            x = layer(x)
        return x

Reload code
import importlib
from feed_forward_neural_network import FeedForwardNeuralNetwork

print((type(nn.Module)))
print(type(FeedForwardNeuralNetwork))
importlib.reload(FeedForwardNeuralNetwork)

Error
&lt;class 'type'&gt;
&lt;class 'type'&gt;
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [53], in &lt;module&gt;
      4 print((type(nn.Module)))
      5 print(type(FeedForwardNeuralNetwork))
----&gt; 6 importlib.reload(FeedForwardNeuralNetwork)

File /usr/lib/python3.9/importlib/__init__.py:140, in reload(module)
    134 &quot;&quot;&quot;Reload the module and return it.
    135 
    136 The module must have been successfully imported before.
    137 
    138 &quot;&quot;&quot;
    139 if not module or not isinstance(module, types.ModuleType):
--&gt; 140     raise TypeError(&quot;reload() argument must be a module&quot;)
    141 try:
    142     name = module.__spec__.name

TypeError: reload() argument must be a module

Environment Details

Python version: 3.9.5
IPython          : 8.0.0
jupyter_client   : 7.1.1
jupyter_core     : 4.9.1


reload() argument should be a module. You need to write it as
importlib.reload(sys.modules.get(FeedForwardNeuralNetwork.__module__))

&lt;module 'feed_forward_neural_network' from '/home/harshit/Downloads/feed_forward_neural_network/feed_forward_neural_network.py'&gt;

