Yeah, I’m afraid that I don’t have much advice for you. It’s been a few months since I looked at the code. The single-TPU fallback would need to be debugged, but I’m not confident what that means or how to do it here. Different versions would be a good way to get things working more easily.