Как преобразовать этот код TensorFlow в PyTorch?

Question

У меня есть этот простой блок кода тензорного потока, что эквивалентно в pytorch? Я застрял, пытаясь закодировать это. Я столкнулся с несколькими ошибками во время выполнения из-за размеров. Это код тензорного потока:

x = tf.placeholder(tf.float32, [708, 256, 3])
y = tf.placeholder(tf.float32, [708, 4])

f1 = tf.Variable(tf.random_normal([5,3,8*3]))
f2 = tf.Variable(tf.random_normal([5,8*3,4*3]))
n1 = tf.Variable(tf.random_normal([61*12,12]))
n2 = tf.Variable(tf.random_normal([12,4]))

b1 = tf.Variable(tf.random_normal([8*3]))
b2 = tf.Variable(tf.random_normal([4*3]))
b3 = tf.Variable(tf.random_normal([12]))
b4 = tf.Variable(tf.random_normal([4]))

conv1 = tf.nn.conv1d(x,f1,stride=1,padding="VALID")
conv1 = tf.nn.bias_add(conv1, b1)
conv1 = tf.nn.sigmoid(conv1)

p1 = tf.layers.average_pooling1d(conv1, pool_size=2, strides=2, padding='VALID')

conv2 = tf.nn.conv1d(p1,f2,stride=1,padding="VALID")
conv2 = tf.nn.bias_add(conv2, b2)
conv2 = tf.nn.sigmoid(conv2)
p2 = tf.layers.average_pooling1d(conv2, pool_size=2, strides=2, padding='VALID')

nn = tf.layers.Flatten()(p2)

fc1 = tf.add(tf.matmul(nn, n1), b3)
fc1 = tf.nn.sigmoid(fc1)
out = tf.add(tf.matmul(fc1, n2), b4)
out = tf.nn.softmax(out)

Как я могу реализовать то же самое в pytorch? Приведенный ниже код - это то, что я пробовал, но размеры перепутаны, я думаю, из-за входов канала.

class TwoLayerNet(torch.nn.Module):
    def __init__(self):
        super(TwoLayerNet,self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1,3, 3*8, kernel_size=5, stride=1),  
            nn.Sigmoid(),
            nn.AvgPool1d(kernel_size=2, stride=0))
        self.conv2 = nn.Sequential(
            nn.Conv1d(3*8, 12, kernel_size=5, stride=1),
            nn.Sigmoid(),
            nn.AvgPool1d(kernel_size=2, stride = 0))
        #self.drop_out = nn.Dropout()

        self.fc1 = nn.Linear(708, 732) 
        self.fc2 = nn.Linear(732, 4)

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = out.reshape(out.size(0), -1)
        out = self.drop_out(out)
        out = self.fc1(out)
        out = self.fc2(out)
        return out