code: def problem_setup(net=LinearNet):
      bf.init()
      num_epochs = 50
      batch_size = 128
      num_train_per_node = 1024
      num_test_per_node = 128
      lr = 0.01
  
      # Setup Problem
      problem_builder = LinearProblemBuilder()
      train_dataset = problem_builder.get_dataset(num_train_per_node)
      train_dataloader = DataLoader(
          train_dataset, batch_size=batch_size
      )
      test_dataset = problem_builder.get_dataset(num_test_per_node)
      test_dataloader = DataLoader(test_dataset, batch_size=batch_size)
      # Setup Model
      model = net(problem_builder.input_dim, problem_builder.output_dim)
      assert (
          num_train_per_node * bf.size() >= model.num_parameters
      ), "The number of samples is too small making it an underdetermined system."
      # Setup Optimizer
      optimizer = optim.Adam(model.parameters(), lr=lr * bf.size())
      bf.broadcast_parameters(model.state_dict(), root_rank=0)
      bf.broadcast_optimizer_state(optimizer, root_rank=0)
      return (
          problem_builder,
          train_dataloader,
          test_dataloader,
          model,
          optimizer,
          num_epochs,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import pytest
  import numpy
  import torch
  import torch.nn
  import torch.optim
  from torch.utils.data import Dataset, DataLoader
  import bluefog.torch
  from bluefog.common import topology_util
label_desc: Use image node_id 1 for calling a global function with example usage: problem_setup(net) and returns: problem_builder,  train_dataloader,  test_dataloader,  model,  optimizer,  num_epochs