Is there any way to use dropout as a function in the forward method? I want to be able to tweak the amount of dropout during training, so something like this would be great:
   layer  = self.fc1(prediction)
   prediction = F.relu(layer)
   prediction = dropout(prediction, prob = parameter) 

where parameter can be learned during training. This would be easy if there was an already written dropout function. With the nn.Dropout, you need to instantiate it with a certain probability, and that remains fixed during training.