like that , sorry i am beginner in cnn:
def forward(self, x):
x.float()
out = self.layer1(x)
out = self.layer2(out)
out = out.reshape(out.size(0), -1)
out = self.drop_out(out)
print(out.shape)
out = self.fc1(out)
out = self.fc2(out)
return out