Running the provided STATE for Virtual Cell Challenge notebook on a T4 Google Colab

[sudhir2016]

The sample notebook provided is marked for A100. But I tried running on T4. During training it just hangs at this step.

Epoch 0: 0% 0/213 [00:00<?, ?it/s]

Waited for 10 minutes. Is it at all possible to run on a T4. And the number of epochs seems to be hardcoded to 213.