How do you create a layer that behaves differently during training and evaluation like dropout?

[shane-carroll]

Do you mean like

def forward(self, x):
    ...
    if self.training:
        x = training_only_layer(x)
    ...