Kerasで書かれたコードをPytorchで実装

前提・実現したいこと

Kerasで書かれたコードをPytorchで実装する際のやり方に困っています。

発生している問題・エラーメッセージ

該当のKerasでかかれたモデルをPytorchで実装しようと考えています。
Pytorch初心者のため、こちらのコードになにか不具合がないかを確認してもらいたいです。
よろしくお願いします。

該当のソースコード

Python
1def loadModel(input_shape):
2 
3   model = Sequential()
4 
5   model.add(Conv2D(16, (5, 5),
6                    activation='relu',
7                    padding='valid',
8                    input_shape=input_shape))
9   model.add(MaxPooling2D((2, 2)))
10   model.add(Conv2D(32, (5, 5),
11                    activation='relu',
12                    padding='valid'))
13   model.add(MaxPooling2D((2, 2)))
14   model.add(Conv2D(64, (5, 5),
15                    activation='relu',
16                    padding='valid'))
17   model.add(Flatten())
18   model.add(Dense(100,
19                   activation='relu'))
20   model.add(Dense(10,
21                   activation='softmax',
22                   use_bias=False))
23 
24   model.compile(loss='categorical_crossentropy',
25                 optimizer='sgd',
26                 metrics=['accuracy'])
27 
28   return model
29

試したこと

Python
1class MyNet(nn.Module):
2    def __init__(self):
3        super(Net, self).__init__()
4        self.relu = nn.ReLU()
5        self.softmax = nn.Softmax(1006)
6
7        self.conv1 = nn.Conv2d(1,16,5)
8        self.conv2 = nn.Conv2d(16,32,5)
9        self.conv3 = nn.Conv2d(32,64,5)
10
11        self.pool = nn.MaxPool2d(2)
12
13        self.flatten = nn.Flatten()
14
15        self.fc1 = nn.Linear(64 * 1 * 1, 100)
16        self.fc2 = nn.Linear(100,10)
17
18    def forward(self, x):
19      x = self.conv1(x)
20      x = self.relu(x)
21      x = self.pool(x)
22      x = self.conv2(x)
23      x = self.relu(x)
24      x = self.pool(x)
25      x = self.conv3(x)
26      x = self.relu(x)
27      x = self.flatten(x)
28      x = self.fc1(x)
29      x = self.fc2(x)
30      return x