def METHOD_NAME(cls):
    weights_before_saving = get_first_weight(cls.learner)
    cls.learner.save(os.path.join(TEMP_SAVE_DIR, 'initial_weights'))
    cls.learner.load('pretrained')
    cls.assertFalse(
        torch.equal(weights_before_saving, get_first_weight(cls.learner)),
        msg="Load() did not alter model weights")
    cls.learner.load(os.path.join(TEMP_SAVE_DIR, 'initial_weights'))
    cls.assertTrue(
        torch.equal(weights_before_saving, get_first_weight(cls.learner)),
        msg="Load did not restore initial weights correctly")
    # Remove temporary files
    try:
        shutil.rmtree(TEMP_SAVE_DIR)
    except OSError as e:
        print(f"Exception when trying to remove temp directory: {e.strerror}")