pytest
torch
numpy