code: def test_torch_prepare_model(ray_start_4_cpus_2_gpus):
      """Tests if ``prepare_model`` correctly wraps in DDP."""
  
      def train_fn():
          model = torch.nn.Linear(1, 1)
  
          # Wrap in DDP.
          model = train.torch.prepare_model(model)
  
          # Make sure model is wrapped in DDP.
          assert isinstance(model, DistributedDataParallel)
  
          # Make sure model is on cuda.
          assert next(model.parameters()).is_cuda
  
      trainer = TorchTrainer(
          train_fn,
          scaling_config=ScalingConfig(num_workers=2, use_gpu=True),
      )
      trainer.fit()
  
      def train_fn_manual_override():
          model = torch.nn.Linear(1, 1)
  
          # Wrap in DDP and manually specify CPU.
          model = train.torch.prepare_model(
              model, device=torch.device("cpu")
          )
  
          # Make sure model is wrapped in DDP.
          assert isinstance(model, DistributedDataParallel)
  
          # Make sure model is NOT on cuda since we manually specified CPU.
          assert not next(model.parameters()).is_cuda
  
      trainer = TorchTrainer(
          train_fn,
          scaling_config=ScalingConfig(num_workers=2, use_gpu=True),
      )
      trainer.fit()
  
imports: import json
  import os
  import time
  from pathlib import Path
  from typing import Dict, List, Union
  from unittest.mock import patch
  import pytest
  import torch
  import torchvision
  from torch.nn.parallel import DistributedDataParallel
  from torch.utils.data import DataLoader, DistributedSampler
  import ray
  import ray.data
  from ray import train
  from ray.exceptions import RayTaskError
  from ray.train import ScalingConfig
  from ray.train._internal.worker_group import WorkerGroup
  from ray.train.constants import DEFAULT_NCCL_SOCKET_IFNAME
  from ray.train.examples.pytorch.torch_linear_example import LinearDataset
  from ray.train.torch.config import TorchConfig, _TorchBackend
  from ray.train.torch.torch_trainer import TorchTrainer
  from ray.train.trainer import TrainingFailedError
label_desc: Use image node_id 5 for calling a global function with example usage: test_torch_prepare_model(ray_start_4_cpus_2_gpus) without return types