torch>=1.6
math
typing-extensions
