Pytorch version

[llllly26]

Hi,

Thanks for your great work! I wonder if this repo could be used to training openVision in a torch version? Since I just find it contains a jax version.

Thanks!