Make Gaussian splatting tests more device-agnostic (#311)

The save/load PLY tests currently assume that `self.device` is the same
as the default `"cuda"` device. In addition, the projection tests assume
that the device being tested is the same as the device stored within the
checkpoint.

Relax both of these assumption in preparation for introducing mGPU
tests. Current test behavior is unchanged.

---------

Signed-off-by: Matthew Cong <mcong@nvidia.com>

Author: matthewdcong

tests/unit/test_gaussian_splat_3d.py

@@ -1150,7 +1150,7 @@ def test_load_ply_with_no_shN(self):
         )
         gs3d_no_shN.save_ply(tf.name)
 
-        gs3d_loaded, metadata = GaussianSplat3d.from_ply(tf.name)
+        gs3d_loaded, metadata = GaussianSplat3d.from_ply(tf.name, device=self.device)
 
         self.assertTrue(torch.allclose(gs3d_loaded.means, gs3d_no_shN.means))
         self.assertTrue(torch.allclose(gs3d_loaded.quats, gs3d_no_shN.quats))
@@ -1247,7 +1247,7 @@ def test_save_and_load_ply(self):
 
         self.gs3d.save_ply(tf.name)
 
-        gs3d_loaded, metadata = GaussianSplat3d.from_ply(tf.name)
+        gs3d_loaded, metadata = GaussianSplat3d.from_ply(tf.name, device=self.device)
 
         self.assertTrue(torch.allclose(gs3d_loaded.means, self.gs3d.means))
         self.assertTrue(torch.allclose(gs3d_loaded.quats, self.gs3d.quats))
@@ -1278,7 +1278,7 @@ def test_save_and_load_ply_with_training_info(self):
         }
         self.gs3d.save_ply(tf.name, metadata=metadata_dict)
 
-        gs3d_loaded, training_info = GaussianSplat3d.from_ply(tf.name)
+        gs3d_loaded, training_info = GaussianSplat3d.from_ply(tf.name, device=self.device)
 
         self.assertTrue(torch.allclose(gs3d_loaded.means, self.gs3d.means))
         self.assertTrue(torch.allclose(gs3d_loaded.quats, self.gs3d.quats))
@@ -1305,7 +1305,7 @@ def test_save_ply_only_string_keys(self):
         metadata_dict = {"_a_key_key": "foo bar baz", "anotherkey": "qux quux corge"}
         self.gs3d.save_ply(tf.name, metadata=metadata_dict)
 
-        gs, meta = GaussianSplat3d.from_ply(tf.name)
+        gs, meta = GaussianSplat3d.from_ply(tf.name, device=self.device)
         self.assertEqual(meta["_a_key_key"], "foo bar baz")
         self.assertEqual(meta["anotherkey"], "qux quux corge")
 
@@ -1315,7 +1315,7 @@ def test_save_ply_only_int_keys(self):
         metadata_dict = {"_a_key_key": 42, "anotherkey": sys.maxsize}
         self.gs3d.save_ply(tf.name, metadata=metadata_dict)
 
-        gs, meta = GaussianSplat3d.from_ply(tf.name)
+        gs, meta = GaussianSplat3d.from_ply(tf.name, device=self.device)
         self.assertEqual(meta["_a_key_key"], 42)
         self.assertEqual(meta["anotherkey"], sys.maxsize)
 
@@ -1353,7 +1353,7 @@ def test_save_and_load_ply_with_training_info_non_contiguous(self):
         }
         self.gs3d.save_ply(tf.name, metadata_dict)
 
-        gs3d_loaded, training_info = GaussianSplat3d.from_ply(tf.name)
+        gs3d_loaded, training_info = GaussianSplat3d.from_ply(tf.name, device=self.device)
 
         self.assertTrue(torch.allclose(gs3d_loaded.means, self.gs3d.means))
         self.assertTrue(torch.allclose(gs3d_loaded.quats, self.gs3d.quats))
@@ -1400,10 +1400,10 @@ def test_gaussian_projection(self):
             torch.save(conics, "regression_conics.pt")
 
         # Regression test
-        test_radii = torch.load(self.data_path / "regression_radii.pt", weights_only=True)
-        test_means2d = torch.load(self.data_path / "regression_means2d.pt", weights_only=True)
-        test_depths = torch.load(self.data_path / "regression_depths.pt", weights_only=True)
-        test_conics = torch.load(self.data_path / "regression_conics.pt", weights_only=True)
+        test_radii = torch.load(self.data_path / "regression_radii.pt", map_location=self.device, weights_only=True)
+        test_means2d = torch.load(self.data_path / "regression_means2d.pt", map_location=self.device, weights_only=True)
+        test_depths = torch.load(self.data_path / "regression_depths.pt", map_location=self.device, weights_only=True)
+        test_conics = torch.load(self.data_path / "regression_conics.pt", map_location=self.device, weights_only=True)
 
         visible = (radii > 0).all(dim=-1)
         torch.testing.assert_close(radii, test_radii)