Update freesplatter/utils/mesh_renderer.py

freesplatter/utils/mesh_renderer.py

@@ -289,14 +289,15 @@ class MeshRenderer(nn.Module):
                  far=10,
                  ssaa=1,
                  texture_filter='linear-mipmap-linear',
-                 opengl=False):
+                 opengl=False,
+                 device='cuda'):
         super().__init__()
         self.near = near
         self.far = far
         assert isinstance(ssaa, int) and ssaa >= 1
         self.ssaa = ssaa
         self.texture_filter = texture_filter
-        self.glctx = dr.RasterizeGLContext() if opengl else dr.RasterizeCudaContext()
+        self.glctx = dr.RasterizeCudaContext(device=device)
 
     def forward(self, meshes, poses, intrinsics, h, w, shading_fun=None,
                 dilate_edges=0, normal_bg=[0.5, 0.5, 1.0], aa=True, render_vc=False):