Update src/pixel3dmm/tracking/flame/lbs.py

src/pixel3dmm/tracking/flame/lbs.py

@@ -208,8 +208,15 @@ def lbs(betas, pose, v_template, shapedirs, posedirs, J_regressor, parents,
     '''
 
     batch_size = max(betas.shape[0], pose.shape[0])
-    device = betas.device
 
+    # after: get the device index (an int), which Dynamo can capture
+    device_idx = betas.get_device()              # returns -1 on CPU, ≥0 for cuda:<idx>
+    if device_idx >= 0:
+        device = torch.device("cuda", device_idx)
+    else:
+        device = torch.device("cpu")
+
+    
     # Add shape contribution
     v_shaped = v_template + blend_shapes(betas, shapedirs)