    def ensure_model_can_train_save_and_load(
        self,
        param_file: Union[PathLike, str],
        tolerance: float = 1e-4,
        cuda_device: int = -1,
        gradients_to_ignore: Set[str] = None,
        overrides: str = "",
        metric_to_check: str = None,
        metric_terminal_value: float = None,
        metric_tolerance: float = 1e-4,
        disable_dropout: bool = True,
        which_loss: str = "loss",
        seed: int = None,
    ):
        """
        # Parameters

        param_file : `str`
            Path to a training configuration file that we will use to train the model for this
            test.
        tolerance : `float`, optional (default=`1e-4`)
            When comparing model predictions between the originally-trained model and the model
            after saving and loading, we will use this tolerance value (passed as `rtol` to
            `numpy.testing.assert_allclose`).
        cuda_device : `int`, optional (default=`-1`)
            The device to run the test on.
        gradients_to_ignore : `Set[str]`, optional (default=`None`)
            This test runs a gradient check to make sure that we're actually computing gradients
            for all of the parameters in the model.  If you really want to ignore certain
            parameters when doing that check, you can pass their names here.  This is not
            recommended unless you're `really` sure you don't need to have non-zero gradients for
            those parameters (e.g., some of the beam search / state machine models have
            infrequently-used parameters that are hard to force the model to use in a small test).
        overrides : `str`, optional (default = `""`)
            A JSON string that we will use to override values in the input parameter file.
        metric_to_check: `str`, optional (default = `None`)
            We may want to automatically perform a check that model reaches given metric when
            training (on validation set, if it is specified). It may be useful in CI, for example.
            You can pass any metric that is in your model returned metrics.
        metric_terminal_value: `str`, optional (default = `None`)
            When you set `metric_to_check`, you need to set the value this metric must converge to
        metric_tolerance: `float`, optional (default=`1e-4`)
            Tolerance to check you model metric against metric terminal value. One can expect some
            variance in model metrics when the training process is highly stochastic.
        disable_dropout : `bool`, optional (default = `True`)
            If True we will set all dropout to 0 before checking gradients. (Otherwise, with small
            datasets, you may get zero gradients because of unlucky dropout.)
        which_loss: `str`, optional (default = `"loss"`)
            Specifies which loss to test. For example, which_loss may be "adversary_loss" for
            `adversarial_bias_mitigator`.
        """
        save_dir = self.TEST_DIR / "save_and_load_test"
        archive_file = save_dir / "model.tar.gz"
        model = train_model_from_file(param_file, save_dir, overrides=overrides, return_model=True)
        assert model is not None

        metrics_file = save_dir / "metrics.json"
        if metric_to_check is not None:
            metrics = json.loads(metrics_file.read_text())
            metric_value = metrics.get(f"best_validation_{metric_to_check}") or metrics.get(
                f"training_{metric_to_check}"
            )
            assert metric_value is not None, f"Cannot find {metric_to_check} in metrics.json file"
            assert metric_terminal_value is not None, "Please specify metric terminal value"
            assert abs(metric_value - metric_terminal_value) < metric_tolerance
        archive = load_archive(archive_file, cuda_device=cuda_device)
        loaded_model = archive.model
        state_keys = model.state_dict().keys()
        loaded_state_keys = loaded_model.state_dict().keys()
        assert state_keys == loaded_state_keys
        # First we make sure that the state dict (the parameters) are the same for both models.
        for key in state_keys:
            assert_allclose(
                model.state_dict()[key].cpu().numpy(),
                loaded_model.state_dict()[key].cpu().numpy(),
                err_msg=key,
            )
        reader = archive.dataset_reader
        params = Params.from_file(param_file, params_overrides=overrides)

        # Need to duplicate params because DataLoader.from_params will consume.
        data_loader_params = params["data_loader"]
        data_loader_params["shuffle"] = False
        data_loader_params2 = Params(copy.deepcopy(data_loader_params.as_dict()))

        if seed is not None:
            random.seed(seed)
            numpy.random.seed(seed)
            torch.manual_seed(seed)

        print("Reading with original model")
        data_loader = DataLoader.from_params(
            params=data_loader_params, reader=reader, data_path=params["validation_data_path"]
        )
        data_loader.index_with(model.vocab)

        if seed is not None:
            random.seed(seed)
            numpy.random.seed(seed)
            torch.manual_seed(seed)

        print("Reading with loaded model")
        data_loader2 = DataLoader.from_params(
            params=data_loader_params2, reader=reader, data_path=params["validation_data_path"]
        )
        data_loader2.index_with(loaded_model.vocab)

        # We'll check that even if we index the dataset with each model separately, we still get
        # the same result out.
        model_batch = next(iter(data_loader))

        loaded_batch = next(iter(data_loader2))

        # Check gradients are None for non-trainable parameters and check that
        # trainable parameters receive some gradient if they are trainable.
        self.check_model_computes_gradients_correctly(
            model, model_batch, gradients_to_ignore, disable_dropout, which_loss
        )

        # The datasets themselves should be identical.
        assert model_batch.keys() == loaded_batch.keys()
        for key in model_batch.keys():
            self.assert_fields_equal(model_batch[key], loaded_batch[key], key, 1e-6)

        # Set eval mode, to turn off things like dropout, then get predictions.
        model.eval()
        loaded_model.eval()
        # Models with stateful RNNs need their states reset to have consistent
        # behavior after loading.
        for model_ in [model, loaded_model]:
            for module in model_.modules():
                if hasattr(module, "stateful") and module.stateful:
                    module.reset_states()
        print("Predicting with original model")
        model_predictions = model(**model_batch)
        print("Predicting with loaded model")
        loaded_model_predictions = loaded_model(**loaded_batch)

        # Both outputs should have the same keys and the values for these keys should be close.
        for key in model_predictions.keys():
            self.assert_fields_equal(
                model_predictions[key], loaded_model_predictions[key], name=key, tolerance=tolerance
            )

        # Check loaded model's loss exists and we can compute gradients, for continuing training.
        loaded_model.train()
        loaded_model_predictions = loaded_model(**loaded_batch)
        loaded_model_loss = loaded_model_predictions[which_loss]
        assert loaded_model_loss is not None
        loaded_model_loss.backward()

        return model, loaded_model
# Parameters

param_file : `str`
Path to a training configuration file that we will use to train the model for this
test.
tolerance : `float`, optional (default=`1e-4`)
When comparing model predictions between the originally-trained model and the model
after saving and loading, we will use this tolerance value (passed as `rtol` to
`numpy.testing.assert_allclose`).
cuda_device : `int`, optional (default=`-1`)
The device to run the test on.
gradients_to_ignore : `Set[str]`, optional (default=`None`)
This test runs a gradient check to make sure that we're actually computing gradients
for all of the parameters in the model.  If you really want to ignore certain
parameters when doing that check, you can pass their names here.  This is not
recommended unless you're `really` sure you don't need to have non-zero gradients for
those parameters (e.g., some of the beam search / state machine models have
infrequently-used parameters that are hard to force the model to use in a small test).
overrides : `str`, optional (default = `""`)
A JSON string that we will use to override values in the input parameter file.
metric_to_check: `str`, optional (default = `None`)
We may want to automatically perform a check that model reaches given metric when
training (on validation set, if it is specified). It may be useful in CI, for example.
You can pass any metric that is in your model returned metrics.
metric_terminal_value: `str`, optional (default = `None`)
When you set `metric_to_check`, you need to set the value this metric must converge to
metric_tolerance: `float`, optional (default=`1e-4`)
Tolerance to check you model metric against metric terminal value. One can expect some
variance in model metrics when the training process is highly stochastic.
disable_dropout : `bool`, optional (default = `True`)
If True we will set all dropout to 0 before checking gradients. (Otherwise, with small
datasets, you may get zero gradients because of unlucky dropout.)
which_loss: `str`, optional (default = `"loss"`)
Specifies which loss to test. For example, which_loss may be "adversary_loss" for
`adversarial_bias_mitigator`.
allennlp/common/testing/model_test_case.py