[Flax] Fix PyTorch import error (#11839)

Commit

4 years ago

[Flax] Fix PyTorch import error (#11839) * fix_torch_device_generate_test * remove @ * change pytorch import to flax import