gemma-2b-it-pytorch on tpu v5p

[shungcp]

Hello, Following the steps in README.md to run gemma-2b-it-pytorch on tpu v5p, i get the error, what's wrong here?

 python scripts/run_xla.py --ckpt /tmp/1/gemma-2b-it.ckpt --variant 2b

TypeError: where() received an invalid combination of arguments - got (Tensor, Tensor, tuple), but expected one of:

• (Tensor condition)
• (Tensor condition, Tensor input, Tensor other, *, Tensor out)
• (Tensor condition, Number self, Tensor other) didn't match because some of the arguments have invalid types: (Tensor, Tensor, tuple of (Tensor, Tensor))
• (Tensor condition, Tensor input, Number other) didn't match because some of the arguments have invalid types: (Tensor, Tensor, tuple of (Tensor, Tensor))
• (Tensor condition, Number self, Number other) didn't match because some of the arguments have invalid types: (Tensor, Tensor, tuple of (Tensor, Tensor))

[michaelmoynihan]

Thanks! I created a PR to address this: https://github.com/google/gemma_pytorch/pull/65