RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _th_bmm #144
Description
🐛 Bugs / Unexpected behaviors
I am trying to put multiple objects into a single mesh, based on example code given in #15 (comment)
I altered the code slightly and have included it below.
When I run the code, it errors out on line 65 with the message:
RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _th_bmm
The full error log is shown at the end of this post.
By the end of the code, everything in the mesh should have been sent to the gpu (RTX 2080 Ti), but for some reason, something is still on the CPU.
Instructions To Reproduce the Issue:
To reproduce, run the script below. I had it saved as "minimal_batch.py"
import torch
from pytorch3d.utils import ico_sphere
from pytorch3d.structures import Textures
from pytorch3d.structures import Meshes
from pytorch3d.renderer import RasterizationSettings
from pytorch3d.renderer import PointLights
from pytorch3d.renderer import look_at_view_transform
from pytorch3d.renderer import OpenGLPerspectiveCameras
from pytorch3d.renderer import MeshRasterizer
from pytorch3d.renderer import MeshRenderer
from pytorch3d.renderer import TexturedSoftPhongShader
device = torch.device("cuda:0")
torch.cuda.set_device(device)
##############################################################################
##############################################################################
# FROM GITHUB
##############################################################################
##############################################################################
mesh1 = ico_sphere(3) # (42 verts, 80 faces)
mesh2 = ico_sphere(4) # (162 verts, 320 faces)
verts1, faces1 = mesh1.get_mesh_verts_faces(0) # Get the vertices and faces
verts2, faces2 = mesh2.get_mesh_verts_faces(0) # Get the vertices and faces
tex1 = torch.ones_like(verts1) # Textures are all ones
tex2 = torch.ones_like(verts2) # Textures are all ones
tex1[:, 1:] *= 0.0 # red
tex2[:, :2] *= 0.0 # blue
verts1 *= 0.25 # Make the red sphere smaller
verts1[:, 0] += 0.8 # Offset first circle
verts2[:, 0] -= 0.5 # Offset second circle
verts = torch.cat([verts1, verts2]) # (204, 3)
faces2 = faces2 + verts1.shape[0] # Offset by the number of vertices in mesh1
faces = torch.cat([faces1, faces2]) # (400, 3)
tex = torch.cat([tex1, tex2])[None] # (1, 204, 3)
textures = Textures(verts_rgb=tex)
##############################################################################
##############################################################################
# END FROM GITHUB
##############################################################################
##############################################################################
mesh = Meshes(verts=[verts], faces=[faces], textures=textures) # Create the mesh
batch_size = 1 # Set batch size
meshes = mesh.extend(batch_size) # Extend the mesh
raster_settings = RasterizationSettings(image_size=256, blur_radius=0.0) # Raster settings
lights = PointLights(device=device, location=[[0.0,0.0,-400.0]]) # Create the lights
elev = torch.linspace(10, 10, batch_size) # Elevation
azim = torch.linspace(-180, 180, batch_size) # Azimuth
R,T = look_at_view_transform(dist=10, elev=elev, azim=azim) # Transformation matrices
cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T) # Create Cameras
rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings) # Rasterizer
shader = TexturedSoftPhongShader(device=device, cameras=cameras, lights=lights) # Shader
renderer = MeshRenderer(rasterizer=rasterizer, shader=shader) # Finally renderer
##############################################################################
##############################################################################
# THROWS ERROR
##############################################################################
##############################################################################
image = renderer(mesh) # (1, H, W, 4)
I ran it with the command "python minimal_batch.py"
Traceback (most recent call last):
File "minimal_batch.py", line 65, in <module>
image = renderer(mesh) # (1, H, W, 4)
File "/usr/lib/python3.8/site-packages/torch/nn/modules/module.py", line 532, in __call__
result = self.forward(*input, **kwargs)
File "/usr/lib/python3.8/site-packages/pytorch3d-0.1.1-py3.8-linux-x86_64.egg/pytorch3d/renderer/mesh/renderer.py", line 51, in forward
fragments = self.rasterizer(meshes_world, **kwargs)
File "/usr/lib/python3.8/site-packages/torch/nn/modules/module.py", line 532, in __call__
result = self.forward(*input, **kwargs)
File "/usr/lib/python3.8/site-packages/pytorch3d-0.1.1-py3.8-linux-x86_64.egg/pytorch3d/renderer/mesh/rasterizer.py", line 113, in forward
meshes_screen = self.transform(meshes_world, **kwargs)
File "/usr/lib/python3.8/site-packages/pytorch3d-0.1.1-py3.8-linux-x86_64.egg/pytorch3d/renderer/mesh/rasterizer.py", line 90, in transform
verts_screen = cameras.transform_points(verts_world, **kwargs)
File "/usr/lib/python3.8/site-packages/pytorch3d-0.1.1-py3.8-linux-x86_64.egg/pytorch3d/renderer/cameras.py", line 244, in transform_points
return world_to_screen_transform.transform_points(points)
File "/usr/lib/python3.8/site-packages/pytorch3d-0.1.1-py3.8-linux-x86_64.egg/pytorch3d/transforms/transform3d.py", line 312, in transform_points
points_out = _broadcast_bmm(points_batch, composed_matrix)
File "/usr/lib/python3.8/site-packages/pytorch3d-0.1.1-py3.8-linux-x86_64.egg/pytorch3d/transforms/transform3d.py", line 680, in _broadcast_bmm
return a.bmm(b)
RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _th_bmm