I trained a model using AlexNet
class OurAlex(nn.Module):
    def __init__(self, num_classes=8):
        super(OurAlex, self).__init__()
        self.alexnet = torchvision.models.alexnet(pretrained=True)
        for param in self.alexnet.parameters():
            param.requires_grad = False

        # Add a avgpool here
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))

        # Replace the classifier layer
        # to customise it according to our output
        self.alexnet.classifier = nn.Sequential(
            nn.Linear(256 * 7 * 7, 1024),
            nn.Linear(1024, 256),
            nn.Linear(256, num_classes),
        )


this model was trained on dogs vs cats dataset
i saved the whole model using
        state = {
            "epoch": epoch,
            "state_dict": model.state_dict(),
            "optimizer": optimizer.state_dict(),
        }
        torch.save(state, 'model.pth')

PROBLEM –
when i tried to load the saved model for predictions and testing the accuracy of it in validation loop
    checkpoint = torch.load("model.pth")
    model.load_state_dict(checkpoint['state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer'])


it gave an error
model.load_state_dict(checkpoint['state_dict'])
TypeError: 'OurAlex' object is not subscriptable

I havent encountered any such error !
Please Help me out in this one …
Thank you
EDIT(adding the train loop)–
model = OurAlex(num_classes=2)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
EPOCHS = 1

TRAIN = False
losses = []


def training_loop(model, optimizer, epochs):
    for epoch in range(epochs):
        try:
            for img, lab in tqdm(Loader):
                img = img.to(device)
                lab = lab.to(device)
                predictions = model(img)
                loss = criterion(predictions, lab)
                optimizer.step()
                optimizer.zero_grad()
                losses.append(loss.item())
                print(f"loss:  {loss.item():.4f}")
        except Exception as e:
            print(str(e))
        state = {
            "epoch": epoch,
            "state_dict": model.state_dict(),
            "optimizer": optimizer.state_dict(),
        }
        torch.save(state, filepath)


if TRAIN:
    training_loop(model, optimizer, EPOCHS)