class cnn_model(nn.Module):
def __init__(self, num_inputs, num_out, activation=nn.ReLU):
super(cnn_model, self).__init__()
self.conv1 = nn.Conv2d(num_inputs, 32, 3, stride=2, padding=1)
self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
self.fc1 = nn.Linear(32 * 6 * 6, 512)
self.fc_out = nn.Linear(512, num_out)
self._initialize_weights()
And the observation space of my env is (4,84,84), 4 is the stack size, so the num_inputs here is 4.