Attribute Error When Saving Pytorch Model

Question

I want to save pytorch model in one .py file and use it in another .py file, but I get AttributeError: 'function' object has no attribute 'copy'

Here is what I have:

train.py

data = {
    "model_state": model.state_dict,
    "optimizer_state": optimizer.state_dict,
    "input_size": input_size,
    "output_size": output_size,
    "hidden_size": hidden_size,
    "all_words": all_words,
    "tags": tags
}
file = "data.pth"
torch.save(data, file)

use.py

file = "data.pth"
data = torch.load(file)
input_size = data['input_size']
output_size = data['output_size']
hidden_size = data['hidden_size']
all_words = data['all_words']
tags = data['tags']
model_state = data['model_state']

model = NeuralNet(input_size, hidden_size, output_size)
model.load_state_dict(model_state)
model.eval()

NeuralNet is a class I defined elsewhere for building a model. The error happens with the second to last line of use.py.