add TPU tests
🚀 Feature
we shall also cover TPU usage as we are supporting it
Motivation
now all changes are tested for GPUs and CPU but we do not have a check for TPU yet
Pitch
getting coverage back to ~99%