code: def test_torch_amp_with_custom_get_state(ray_start_4_cpus):
      """Tests amp with a model that has a custom __getstate__ method defined.
  
      See https://discuss.ray.io/t/ray-train-hangs-for-long-time/6333/7
      """
  
      def train_fn():
          train.torch.accelerate(amp=True)
  
          class CustomLinear(torch.nn.Linear):
              def __getstate__(self):
                  return self.__dict__.copy()
  
          model = CustomLinear(1, 1)
          model = train.torch.prepare_model(model)
  
          # TorchCheckpoint.from_model fails, so just save the state dict only.
          train.report(
              {},
              checkpoint=TorchCheckpoint.from_state_dict(
                  model.module.state_dict()
              ),
          )
  
      trainer = TorchTrainer(
          train_fn,
          scaling_config=ScalingConfig(num_workers=2),
      )
      results = trainer.fit()
      assert results.checkpoint
  
imports: import contextlib
  import os
  import tempfile
  import time
  import uuid
  from unittest.mock import patch
  import pytest
  import torch
  import ray
  import ray.train
  from ray.cluster_utils import Cluster
  from ray.train import Checkpoint, RunConfig, ScalingConfig
  from ray.train.examples.pytorch.torch_linear_example import train_func
  from ray.train.torch import TorchCheckpoint, TorchConfig, TorchPredictor, TorchTrainer
  from ray.train.trainer import TrainingFailedError
label_desc: Use image node_id 12 for calling a global function with example usage: test_torch_amp_with_custom_get_state(ray_start_4_cpus) and returns: self