xla_fit_one_cycle fails with an error when using fastai version 2.3.1

[butchland]

When running an xla fit method, it triggers an error:

Exception in device=TPU:0: Input type (torch.FloatTensor) and weight type (XLAFloatType) should be the same

Note: downgrading to 2.3.0 is a temporary workaround while this is being fixed Works with fastai==2.3.0 See this gist for a error sample