torch.bernoulli expects an input tensor containing the probabilities of drawing a 1 value. So depending how p is defined your code should be correct.
Note that you would have to take care of the scaling in dropout layers (either the inverse scaling during training or the vanilla scaling during evaluation).
To switch the behavior between training and evaluation you could create a custom nn.Module and use the internal self.training flag to switch between the behaviors.
The self.training flag will be changed through calls into model.train() and model.eval().