cerebras.modelzoo.common.run_cstorch_flow.run_cstorch_flow#

cerebras.modelzoo.common.run_cstorch_flow.run_cstorch_flow(params, model_fn, train_data_fn, eval_data_fn)[source]#

Set up the cstorch run and call the appropriate helper based on the mode

Parameters
  • params – the params dictionary extracted from the params.yaml used

  • model_fn – A callable that takes in the params dictionary and returns a torch.nn.Module

  • train_data_fn – A callable that takes in the param dictionary and returns a torch.utils.data.DataLoader

  • eval_data_fn – A callable that takes in the param dictionary and returns a torch.utils.data.DataLoader