I saved my model with this code:

from google.colab import files
torch.save(net, 'model.pth')



Then uploaded this way and checked on an image (x):

model = torch.load('model.pth')
model.eval()
torch.argmax(model(x))


And on the old session, it worked great, but then I started a new session and tried the code above, and got such an error:

AttributeError: Can't get attribute 'EfficientNet' on <module '__main__'>


Maybe somebody knows how to deal with it?

Got the very same error recently.

Your network is usually defined as a class (here class EfficientNet(nn.Module). It seems when we load a model, it needs the class to be defined so it can instantiate it.

In my case, the class was defined in the training .py file. So what I did to fix that error was just copy-paste (it seems importing it didn't work for me, so I had to copy paste) the whole class definition to the file where you load your model.

In my case the class definition looked something like that and needed to be pasted before using the torch.load function:

class FCNClass(nn.Module):
def __init__(self, in_feat=1000, nb_classes=3, nb_hid=1, hidden_size=500, act=nn.ReLU):
super(FCNClass, self).__init__()

self.act = act()
self.flat = nn.Flatten()

self.fc1 = nn.Linear(in_feat, hidden_size)
self.fcs = nn.ModuleList()
for nb in range(nb_hid):
self.out = nn.Linear(hidden_size, nb_classes)

def forward(self, x):
x = self.flat(x)
x = self.fc1(x)
x = self.act(x)
for lay in self.fcs:
x = lay(x)
x = self.act(x)
x = self.out(x)
x = self.act(x)
return x


The class basically defines the architecture of the model, if it is not defined, then the load function doesn't know where to put the weights. Mine was done with torch.load(path) and not model.load_state_dict(torch.load(path)) but it should not make a difference
• post edited, hopes it clarify your understanding, pytorch.org/tutorials/beginner/…. Here is the tutorial for loading models with state dictionnary. Maybe the issue in your code is when you save your network, be sure to save the state_dict if you load the state_dict : torch.save(model.state_dict(), PATH) May 2, 2021 at 16:51
You need to first define the model. Once you have defined the model, then, instantiate a class of it. Once that is done, use model.load_state_dict(torch.load(path_to_model_file)).