def METHOD_NAME(self):
    model = RegressionModel()
    # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
    distributed_model = torch.nn.parallel.DataParallel(model)
    model_unwrapped = extract_model_from_parallel(distributed_model)
    self.assertEqual(model, model_unwrapped)