code: def test_torch_fsdp(ray_start_4_cpus_2_gpus):
      """Tests if ``prepare_model`` correctly wraps in FSDP."""
  
      def train_fn():
          model = torch.nn.Linear(1, 1)
  
          # Wrap in FSDP.
          model = train.torch.prepare_model(
              model, parallel_strategy="fsdp"
          )
  
          # Make sure model is wrapped in FSDP.
          assert isinstance(model, FullyShardedDataParallel)
  
          # Make sure the 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()
  
file_path: ray/python/ray/train/tests/test_torch_fsdp.py
imports: import pytest
  import torch
  from torch.distributed.fsdp import FullyShardedDataParallel
  import ray
  from ray import train
  from ray.train import ScalingConfig
  from ray.train.torch import TorchTrainer
label_desc: Use image node_id 2 for calling a global function with example usage: test_torch_fsdp(ray_start_4_cpus_2_gpus) without return types