From 08f5ff787f73f1d8acd9dce3331caa6e273b0670 Mon Sep 17 00:00:00 2001
From: Peder Bergebakken Sundt <pbsds@hotmail.com>
Date: Fri, 25 Oct 2024 16:37:04 +0000
Subject: [PATCH] python312Packages.pytorch3d: test cuda

---
 .../python-modules/pytorch3d/default.nix            | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/pkgs/development/python-modules/pytorch3d/default.nix b/pkgs/development/python-modules/pytorch3d/default.nix
index 1a9151591863..2a92a51ec8bf 100644
--- a/pkgs/development/python-modules/pytorch3d/default.nix
+++ b/pkgs/development/python-modules/pytorch3d/default.nix
@@ -46,6 +46,19 @@ buildPythonPackage rec {
 
   pythonImportsCheck = [ "pytorch3d" ];
 
+  passthru.tests.rotations-cuda =
+    cudaPackages.writeGpuTestPython { libraries = ps: [ ps.pytorch3d ]; }
+      ''
+        import pytorch3d.transforms as p3dt
+
+        M = p3dt.random_rotations(n=10, device="cuda")
+        assert "cuda" in M.device.type
+        angles = p3dt.matrix_to_euler_angles(M, "XYZ")
+        assert "cuda" in angles.device.type
+        assert angles.shape == (10, 3), angles.shape
+        print(angles)
+      '';
+
   meta = {
     description = "FAIR's library of reusable components for deep learning with 3D data";
     homepage = "https://github.com/facebookresearch/pytorch3d";