Hi @IamSparky, sorry for the late reply, we take all xla/tpu related questions/feature requests/bug reports as issues in pytorch/xla repo. Would you mind posting the issue in pytorch/xla github with a link to the colab notebook? I’m pretty sure this should work but it’s hard to locate without seeing the code.