code: def test_layerwise_sandbox(device):
      model = nn.Sequential(nn.Conv2d(3, 3, 1), nn.BatchNorm2d(3))
      model.eval()
  
      for layer in layerwise_sandbox(model, torch.device(device)):
          assert layer.training
          assert all(
              p.device.type == device for p in layer.parameters()
          )
  
      assert all(not l.training for l in model)
      assert all(p.device.type == "cpu" for p in model.parameters())
  
imports: import time
  import pytest
  import torch
  from torch import nn
  from torch.distributed.pipeline.sync._balance import balance_by_size, balance_by_time, blockpartition
  from torch.distributed.pipeline.sync._balance.profile import layerwise_sandbox
  from torch.testing._internal.common_utils import run_tests
label_desc: Use image node_id 10 for calling a global function with example usage: test_layerwise_sandbox(device) without return types