code: def _train_reg_model(model, device, offload_device, use_fp16=False):
      reg_model = copy.deepcopy(model)
      reg_model = reg_model.cuda()
      reg_optimizer = torch.optim.SGD(reg_model.parameters(), lr=0.001)
      return _train(reg_model, reg_optimizer, use_fp16, device)
  
file_path: fairscale/tests/experimental/nn/test_offload.py
imports: import contextlib
  import copy
  import numpy
  import pytest
  import torch
  from fairscale.experimental.nn.offload import OffloadModel
  from fairscale.fair_dev.testing.testing import skip_if_no_cuda
  from fairscale.internal import torch_version
label_desc: Use image node_id 7 for calling a global function with example usage: _train_reg_model(model, device, offload_device, use_fp16) and returns: _train