Could you try following what I said above?
So, save the model as,
torch.save({'model_state_dict': model.state_dict()}, "model.pt")

and try loading it back with,
state_dict = torch.load(f="model.pt", map_location=device) #device is torch.device('cpu') or torch.device("cuda")
model.load_state_dict(state_dict['model_state_dict'])

This is slightly more convoluted but it might work! If this fails, could you print out the keys of the Dictionary?
for key, value in state_dict:
  print(key)