The pytorch training model cannot be created successfully
I would like to do a neural network for regression analysis using optuna based on this site.
I would like to create a model with two 1D data as input and one 1D data as output in batch learning.
x is the training data and y is the teacher data.
class Model(nn.Module):
    # コンストラクタ(インスタンス生成時の初期化)
    def __init__(self,trial, mid_units1, mid_units2):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(2, mid_units1)
        self.bn1 = nn.BatchNorm1d(mid_units1)
        self.linear2 = nn.Linear(mid_units1, mid_units2)
        self.bn2 = nn.BatchNorm1d(mid_units2)
        self.linear3 = nn.Linear(mid_units2, 1)
        self.activation = trial_activation(trial)

    def forward(self, x):
        x = self.linear1(x)
        x = self.bn1(x)
        x = self.activation(x)
        x = self.linear2(x)

device = &quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;

EPOCH = 100
x = torch.from_numpy(a[0].astype(np.float32)).to(device)
y = torch.from_numpy(a[1].astype(np.float32)).to(device)

def train_epoch(model, optimizer, criterion):
    model.train()
    optimizer.zero_grad()    # 勾配情報を0に初期化
    y_pred = model(x)                                               # 予測
    loss = criterion(y_pred.reshape(y.shape), y)          # 損失を計算(shapeを揃える)
    loss.backward()                                                       # 勾配の計算
    optimizer.step()                                                      # 勾配の更新
    return loss.item()

def trial_activation(trial):
    activation_names = ['ReLU','logsigmoid']
    activation_name = trial.suggest_categorical('activation', activation_names)
    if activation_name == activation_names[0]:
        activation = F.relu
    else:
        activation = F.logsigmoid
    return activation

def objective(trial):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    # 中間層のユニット数の試行
    mid_units1 = int(trial.suggest_discrete_uniform(&quot;mid_units1&quot;, 1024*2,1024*4, 64*2))
    mid_units2 = int(trial.suggest_discrete_uniform(&quot;mid_units2&quot;, 1024, 1024*2, 64*2))

    net = Model(trial, mid_units1, mid_units2).to(device)

    criterion = nn.MSELoss() 
    # 最適化手法の試行
    optimizer = trial_optimizer(trial, net)
    train_loss = 0
    for epoch in range(EPOCH):
        train_loss = train_epoch(net, optimizer, criterion, device)
    torch.save(net.state_dict(), str(trial.number) + &quot;new1.pth&quot;)
    return train_loss

strage_name = &quot;a.sql&quot;
study_name = 'a'

study = optuna.create_study(
    study_name = study_name,
    storage='sqlite:///'  + strage_name, 
    load_if_exists=True,
    direction='minimize')
TRIAL_SIZE = 100

study.optimize(objective, n_trials=TRIAL_SIZE)

error message
---&gt; 28     loss = criterion(y_pred.reshape(y.shape), y)          # 損失を計算(shapeを揃える)
     29     loss.backward()                                                       # 勾配の計算
     30     optimizer.step()                                                      # 勾配の更新

AttributeError: 'NoneType' object has no attribute 'reshape'

Because of the above error, I checked the value of y_pred and found it to be None.
    model.train()
    optimizer.zero_grad()

I am thinking that these two lines may be wrong, but I don't know how to solve this problem.

With PyTorch, when you call y_pred = model(x) that will call the forward function which is defined in the Model class.
So, y_pred  will get the result of the forward function, in your case, it returns nothing, that's why you get a None value. You can change the forward function as below:
    def forward(self, x):
        x = self.linear1(x)
        x = self.bn1(x)
        x = self.activation(x)
        x = self.linear2(x)
        return x

