Function create_model

tests/unit/runtime/test_multiple_models.py:14–18

(config_dict)

def create_model(config_dict):
    hidden_dim = 64
    model = SimpleModel(hidden_dim)
    model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
    return model


def train_shared_loss(num_models, config_dict, dtype):