pyrox.dev
lol
python312Packages.pytorch3d: test cuda
authored by
Peder Bergebakken Sundt
and committed by
SomeoneSerge
08f5ff78
030fc6a5
+13
pkgs/development/python-modules/pytorch3d/default.nix
+13
pkgs/development/python-modules/pytorch3d/default.nix
···
46
46
47
47
pythonImportsCheck = [ "pytorch3d" ];
48
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
+
49
62
meta = {
50
63
description = "FAIR's library of reusable components for deep learning with 3D data";
51
64
homepage = "https://github.com/facebookresearch/pytorch3d";