accepts lists to save meshes

neural_monitor/monitor.py

@@ -1222,19 +1222,17 @@ def save_meshes(self, name, meshes=None, verts=None, faces=None, verts_uvs=None,
                 '`meshes` and other arguments are mutually exclusive'
 
         if meshes is None:
-            if len(verts.shape) == len(faces.shape) == 2:
-                filename = os.path.join(self.mesh_folder, f'{name}.obj')
-                io.save_obj(filename, verts, faces, verts_uvs=verts_uvs, faces_uvs=faces_uvs, texture_map=texture_map)
-                return
-
-            if len(verts.shape) == len(faces.shape) == 3:
-                assert verts.shape[0] == faces.shape[0]
-                if verts_uvs is not None:
-                    assert verts_uvs.shape[0] == verts.shape[0] == faces_uvs.shape[0] == texture_map.shape[0]
-                else:
-                    verts_uvs = faces_uvs = texture_map = [None] * verts.shape[0]
+            if isinstance(verts, T.Tensor):
+                if len(verts.shape) == len(faces.shape) == 2:
+                    filename = os.path.join(self.mesh_folder, f'{name}.obj')
+                    io.save_obj(filename, verts, faces, verts_uvs=verts_uvs, faces_uvs=faces_uvs, texture_map=texture_map)
+                    return
+
+            assert len(verts) == len(faces)
+            if verts_uvs is not None:
+                assert len(verts_uvs) == len(verts) == len(faces_uvs) == len(texture_map)
             else:
-                raise NotImplementedError
+                verts_uvs = faces_uvs = texture_map = [None] * len(verts)
         else:
             verts = meshes.verts_list()
             faces = meshes.faces_list()