mirror of
https://github.com/vale981/poetry2nix
synced 2025-03-06 09:41:39 -05:00
15 lines
711 B
Nix
15 lines
711 B
Nix
{ poetry2nix, python310, runCommand }:
|
|
let
|
|
env = poetry2nix.mkPoetryEnv {
|
|
python = python310;
|
|
pyproject = ./pyproject.toml;
|
|
poetrylock = ./poetry.lock;
|
|
preferWheels = true;
|
|
};
|
|
in
|
|
# Note: torch.cuda() will print False, even if you have a GPU, when this runs *during test.*
|
|
# But if you run the script below in your shell (rather than during build), it will print True.
|
|
# Presumably this is due to sandboxing.
|
|
runCommand "ml-stack-test" { } ''
|
|
${env}/bin/python -c 'import torch; import torchvision; print(torch.__version__); print(torchvision.__version__); print(torch.cuda.is_available()); x = torch.randn(3,3); x = x.cuda() if torch.cuda.is_available() else x; print(x**2)' > $out
|
|
''
|