def METHOD_NAME(
    move_to_device: Callable,
    precision_context,
    input_dtype=torch.float32,
):
    make_deterministic(warn_only=True)
    memory_stats = {}
    model = ConvNet()
    model = move_to_device(model)
    dataloader = model.get_dataloader()
    optimizer = model.get_optimizer()
    loss_fn = model.get_loss_function()
    memory_stats["start"] = torch.cuda.memory_stats()
    model.train()
    iteration_timings = []
    iterator = iter(dataloader)
    for _ in range(model.num_steps):
        t0 = time.perf_counter()
        inputs, labels = next(iterator)
        inputs, labels = move_to_device(inputs), move_to_device(labels)
        optimizer.zero_grad()
        with precision_context():
            outputs = model(inputs.to(input_dtype))
        loss = loss_fn(outputs.float(), labels)
        loss.backward()
        optimizer.step()
        t1 = time.perf_counter()
        iteration_timings.append(t1 - t0)
    memory_stats["end"] = torch.cuda.memory_stats()
    return model.state_dict(), torch.tensor(iteration_timings), memory_stats