if self.track_running_stats:
            self.register_buffer('running_mean', torch.zeros(num_features))
            self.register_buffer('running_var', torch.ones(num_features))
            self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
        else:
            self.register_parameter('running_mean', None)
            self.register_parameter('running_var', None)
            self.register_parameter('num_batches_tracked', None)

@albanD, a nitpicking: why running_mean, running_var and num_batches_tracked in case of track_running_stats=False are registered as parameters ? I understand that this have no effect but this looks a bit strange…