I tried to run your code on a google cloud vm but I did not see the error you mentioned( I have to remove flags argument in your spawn_function since you pass args=() ). This kind of problem usually caused by your code calling XLA functions at global scope, together with multi-processing, which is not allowed. I am guessing your did that prior to this code and the machine is being left in a bad state? Restarting the kaggle kernel might help.
Issues · pytorch/xla · GitHub 13 is a better place to ask pt/xla related question. If you can share the kaggle notebook we might be able to help you to see what is going wrong.