torch>=1.6
jax>=0.2.20
