Pytorch: Disable only nn.Dropout() without using model.eval()

Question

nn.Dropout() can be disabled by using model.eval().
However by using .eval(), nn.BatchNorm1d() are also disabled. Because the distributions between train and test sets are different, I'd like to disable only Dropout for generating data by GAN.
Is there any way to disable only Dropout after training?
Here is the generator model in my GAN.

class Generator(nn.Module):
def __init__(self, num_input=2, noise_dim=1, num_output=5, hidden_size=128):
    super(Generator, self).__init__()

    self.fc_in = nn.Linear(num_input+noise_dim, hidden_size)
    self.fc_mid = nn.Linear(hidden_size+num_input+noise_dim, hidden_size)
    self.fc_out = nn.Linear(2*hidden_size+num_input+noise_dim, num_output)

    self.bn_in = nn.BatchNorm1d(hidden_size)
    self.bn_mid = nn.BatchNorm1d(hidden_size)
    
    self.dropout = nn.Dropout()
    
    self.relu = nn.ReLU()

def forward(self, y, z):
    h0 = torch.concat([y,z],axis=1)
    h1 = self.relu(self.bn_in(self.fc_in(h0)))
    h1 = self.dropout(h1)
    h1 = torch.concat([h0,h1],axis=1)
    h2 = self.relu(self.bn_mid(self.fc_mid(h1)))
    h2 = self.dropout(h2)
    h2 = torch.concat([h1,h2],axis=1)
    x = self.fc_out(h2)
    return x

Answer 1

The answer in the comment is right, you can run eval() on single modules.

But... why do you think you need to keep the BatchNorm active after training? By default, in eval mode it will use the running average/std computed during training (which is a good thing, and makes the model give the same outputs regardless of the batch size or the content of the batch). The outputs should be closer to the ones in training when the batchnorm is in eval mode. Sometimes it's beneficial to recompute the batch norm mean/std after the training (but still later using the whole model in eval mode afterwards). There are some rare cases when you might want to use it in train mode, but make sure that it's what you want.