TPU support for gossipy about gossipy HOT 2 OPEN

Hi @ParsaMokhtariHessar , thank you for your interest in gossipy. It seems that what you have done is correct, so it's hard to say what is wrong. Have you tried working with TPUs using a neural net outside gossipy? So to be sure that the problem comes from the framework. Anyway, (this happens also with GPUs), in these simulations there are a bunch of models and only one of them is loaded into the GPU/TPU memory at a time, and what happens is that the time overhead to move the models in/out of the memory usually overcomes the benefits of using GPU/TPUs. This is my take on this, but I honestly have to say that I did not dig into it.

from gossipy.

I see! And I guess it would be too much memory to attempt to move all calculations to the TPU VRAM. Nonetheless the "Loading trip" is significantly faster in the case of a GPU. I am going to attempt to time a simple MNIST train with CPU, GPU, and TPU and get back to you.

from gossipy.