PyTorch and Pickle Error: AttributeError: Can't get attribute "Class_name" on

I have created a neural network model using PyTorch for a time-series forecasting problem. I have saved the model using Pickle.When loading the model to check on test data, it is throwing an attribute error.Below, I'm providing the relevant parts of my code:

import torchimport torch.nn as nnimport torch.nn.functional as Fclass twoD_predict(nn.Module):  def __init__(self):    super().__init__()  def forward(self,x):  ...  def train(self,        epochs = 100):  ...obj1 = twoD_predict()obj1.train()

I saved the model using Pickle as follows:

import picklefilename = (f"{column_names[0]}.sav")pickle.dump(obj1, open(filename, 'wb'))

However, when I try to load the model with the following code:

import pickleif __name__ == '__main__':    with open("model.sav", 'rb') as file:        model = pickle.load(file)

I encounter the error: "AttributeError: Can't get attribute "twoD_predict" on <module '__main__' from '/load_model.py'>"

Can anyone please help?