code: def _get_model(num_inputs=2, num_hidden=20, num_outputs=2):
      num_layers = torch.cuda.device_count() - 2
      model = torch.nn.Sequential(
          torch.nn.Linear(num_inputs, num_hidden),
          *(
              [
                  torch.nn.Linear(num_hidden, num_hidden)
                  for _ in range(num_layers)
              ]
          ),
          torch.nn.Linear(num_hidden, num_outputs),
      )
      return model
  
file_path: fairscale/tests/nn/pipe/test_parity.py
imports: import contextlib
  import copy
  import numpy
  import pytest
  import torch
  from fairscale.fair_dev.testing.testing import skip_if_single_gpu
  from fairscale.nn import Pipe
label_desc: Use image node_id 1 for calling a global function with example usage: _get_model(num_inputs, num_hidden, num_outputs) and returns: model