pyrox.dev / nixpkgs
lol
fork atom

python312Packages.pytorch3d: test cuda

authored by

Peder Bergebakken Sundt and committed by
SomeoneSerge 08f5ff78 030fc6a5

+13
unified split
+13
pkgs/development/python-modules/pytorch3d/default.nix
··· 46 47 pythonImportsCheck = [ "pytorch3d" ]; 48 49 meta = { 50 description = "FAIR's library of reusable components for deep learning with 3D data"; 51 homepage = "https://github.com/facebookresearch/pytorch3d";
··· 46 47 pythonImportsCheck = [ "pytorch3d" ]; 48 49 + passthru.tests.rotations-cuda = 50 + cudaPackages.writeGpuTestPython { libraries = ps: [ ps.pytorch3d ]; } 51 + '' 52 + import pytorch3d.transforms as p3dt 53 + 54 + M = p3dt.random_rotations(n=10, device="cuda") 55 + assert "cuda" in M.device.type 56 + angles = p3dt.matrix_to_euler_angles(M, "XYZ") 57 + assert "cuda" in angles.device.type 58 + assert angles.shape == (10, 3), angles.shape 59 + print(angles) 60 + ''; 61 + 62 meta = { 63 description = "FAIR's library of reusable components for deep learning with 3D data"; 64 homepage = "https://github.com/facebookresearch/pytorch3d";