-record(train_config, {
    num_workers :: integer(),
    batches_per_step :: integer(),
    grad_aggregation :: viva_tensor@distributed@trainer:grad_aggregation()
}).