code: def input_singleton(pipe_class):
      class One(nn.Module):
          def __init__(self):
              super().__init__()
              self.fc = nn.Linear(1, 1)
  
          def forward(self, only_a):
              (a,) = only_a
              return (self.fc(a),)
  
      model = nn.Sequential(One())
      model = pipe_class(
          model,
          balance=[1],
          worker_map=get_worker_map(),
          chunks=2,
      )
  
      a = torch.rand(10, 1, requires_grad=True)
  
      (a_out,) = model((a,))
      loss = a_out.mean()
      loss.backward()
  
      assert all(p.grad is not None for p in model.parameters())
      assert a.grad is not None
  
imports: from collections import OrderedDict
  from copy import deepcopy
  import os
  import time
  import pytest
  import torch
  from torch import nn
  from fairscale.fair_dev.testing.testing import get_worker_map, torch_spawn
  from fairscale.internal import torch_version
  from fairscale.nn.model_parallel.initialize import get_pipeline_parallel_group
  from fairscale.nn.pipe import AsyncPipe
  from fairscale.nn.pipe.types import LazyModule
label_desc: Use image node_id 23 for calling a global function with example usage: input_singleton(pipe_class) and returns: self