Function generate_data

intermediate_source/torch_compile_full_example.py:80–84

(b)

# Generates random input and targets data for the model, where `b` is
# batch size.
def generate_data(b):
    return (
        torch.randn(b, 3, 128, 128).cuda(),
        torch.randint(1000, (b,)).cuda(),
    )


N_ITERS = 10