code: def test_checkpoint_torch_model_with_amp(ray_start_4_cpus_2_gpus):
      """Test that model with AMP is serializable."""
  
      def train_func():
          train.torch.accelerate(amp=True)
  
          model = torchvision.models.resnet101()
          model = train.torch.prepare_model(model)
  
          with TemporaryDirectory() as tmpdir:
              torch.save(model, os.path.join(tmpdir, "checkpoint.pt"))
              train.report(
                  {}, checkpoint=Checkpoint.from_directory(tmpdir)
              )
  
      trainer = TorchTrainer(
          train_func,
          scaling_config=ScalingConfig(num_workers=2, use_gpu=True),
      )
      results = trainer.fit()
      assert results.checkpoint
  
file_path: ray/python/ray/train/tests/test_gpu_amp.py
imports: import os
  from tempfile import TemporaryDirectory
  from timeit import default_timer
  import torch
  import torchvision
  import ray.train
  from ray.train import Checkpoint, ScalingConfig
  from ray.train.torch.torch_trainer import TorchTrainer
label_desc: Use image node_id 2 for calling a global function with example usage: test_checkpoint_torch_model_with_amp(ray_start_4_cpus_2_gpus) without return types