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

src/pixel3dmm/tracking/flame/lbs.py

@@ -210,12 +210,7 @@ def lbs(betas, pose, v_template, shapedirs, posedirs, J_regressor, parents,
     batch_size = max(betas.shape[0], pose.shape[0])
 
     # 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")
-
+    device = betas.device
     
     # Add shape contribution
     v_shaped = v_template + blend_shapes(betas, shapedirs)