def METHOD_NAME(self, bn=False):
    model_new = get_model_from_arch_str(self.arch_str, self.num_classes, use_bn=bn)
    # TODO this is quite brittle and doesn't work with nn.Sequential when bn is different
    # it is only required to maintain initialization -- maybe init after get_punable_copy?
    model_new.load_state_dict(self.state_dict(), strict=False)
    model_new.train()
    return model_new