WARNING: PROCEED WITH CAUTION. WHILE  OVERRIDING NN.MODULE.TRAIN() CAN BE HELPFUL TO CONTROL THE FROZEN PARAMETERS, IT CAN BACKFIRE IF YOU FORGET SELF.FREEZE SET TO TRUE WHILE YOU INTEND TO TRAIN THE ENTIRE MODEL.
as @EthanZhangYi did, i recommend to override your  model.train() because when you call the original, batchnorm will be back on (the running avg, var [buffers only], but not the params). it is safer this way.
    def train(self, mode=True):
        """
        Override nn.Module.train() to consider the freezing aspect if required.
        """
        super(YourModel, self).train(mode=mode)  # will turn on batchnorm (buffers not params).

        if self.freeze:
            self.freeze_stuff()   # call freeze to turn off the batch-norm. 

        return self

self.freeze_stuff() should take care of parameters, batchnorm, and dropout (at least).