diff --git a/smpl/smpl_torch.py b/smpl/smpl_torch.py
index fbac7fb..1b90f78 100644
--- a/smpl/smpl_torch.py
+++ b/smpl/smpl_torch.py
@@ -163,9 +163,12 @@ class SMPL(Module):
     v_shaped = torch.stack(v_shaped, axis=-1)
     
     # J = torch.matmul(self.J_regressor, v_shaped)
-    J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])])
-    J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
-    J = J - J0
+    if verts is None:
+      J = torch.stack([torch.matmul(self.J_regressor[g.to(torch.long)], v_shaped[i, :, :, g.to(torch.long)]) for i, g in enumerate(betas[:,-1])])
+    else:
+      J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])])
+      J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
+      J = J - J0
     R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
 
     if verts is not None:
@@ -208,7 +211,8 @@ class SMPL(Module):
     
     joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
             
-    T = torch.stack([torch.tensordot(results[i], self.weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
+    weights = self.weights if verts is not None else self.weights_all 
+    T = torch.stack([torch.tensordot(results[i], weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])]) 
     rest_shape_h = torch.cat(
       (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
     )