[tests] add test for group_offloading with training.

[sayakpaul]

What does this PR do?

@a-r-r-o-w currently it has:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and CPU!

When run with:

pytest tests/models/transformers/test_models_transformer_flux.py::FluxTransformerTests::test_group_offloading_with_training_0

Is this known?

[sayakpaul]

@a-r-r-o-w a gentle ping.