diff --git a/custom_utils/smpl_2.py b/custom_utils/smpl_2.py
index af2deac..bc12542 100644
--- a/custom_utils/smpl_2.py
+++ b/custom_utils/smpl_2.py
@@ -22,7 +22,7 @@ eval_joints = {
         
 class SMPL:
 
-    def __init__(self):
+    def __init__(self, path):
         self.J_regressor = []
         self.v_template = []
         self.shapedirs = []
@@ -30,7 +30,7 @@ class SMPL:
         self.weights = []
         # Load SMPL models
         # path = os.path.abspath(os.path.dirname(__file__))
-        with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f:
+        with open(os.path.join(path, 'super_model_f.pkl'), 'rb') as f:
             if sys.version_info[0]<3:
                 female = pickle.load(f)
             else:
@@ -40,7 +40,7 @@ class SMPL:
             self.posedirs.append(tf.constant(female['posedirs'], dtype=np.float32))
             self.shapedirs.append(tf.constant(female['shapedirs'], dtype=np.float32))
             self.weights.append(tf.constant(female['weights'], dtype=np.float32))
-        with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f:
+        with open(os.path.join(path, 'super_model_m.pkl'), 'rb') as f:
             if sys.version_info[0]<3:
                 male = pickle.load(f)
             else:
@@ -63,11 +63,11 @@ class SMPL:
         self.weights_one_hot = tf.one_hot(self.weights_one_hot, 24)
 
         # Topology no head/hands/feet
-        F = sio.loadmat('custom_utils/ssmpl_faces.mat')
+        F = sio.loadmat(path + '/ssmpl_faces.mat')
         self.f = np.array(F['faces_default'])
         self.f_skirt = np.array(F['faces_skirt'])
-        self.nn = np.array(sio.loadmat('custom_utils/NN.mat'))
-        self.to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
+        self.nn = np.array(sio.loadmat(path + '/NN.mat'))
+        self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1)
 
         self.weights = tf.gather(self.weights, self.to_keep, axis=1)
         # self.shapedirs = tf.gather(self.shapedirs, self.to_keep, axis=1)
@@ -246,14 +246,13 @@ class SMPL:
         return J
 
 
-    def run_garment(self, pose, betas, gender, v_cloth, filter=None):
+    def run_garment(self, R, betas, gender, v_cloth, filter_idx=None):
         batch_size = v_cloth.shape[0]
 
         J = self.get_joints_template(betas, gender, batch_size)
-        R = self.run_rodrigues(pose)
         
-        if filter is not None:
-            weights = tf.boolean_mask(self.weights, filter, axis=1)
+        if filter_idx is not None:
+            weights = tf.gather(self.weights, filter_idx, axis=1)
 
         R_cube_big = R
         results = []
@@ -274,6 +273,8 @@ class SMPL:
         v = tf.reshape(v, (batch_size, -1, 4))[:, :, :3]
 
         J_out = tf.squeeze(J)
+        if J.shape[0] == 1:
+            J_out = J_out[tf.newaxis, :, :]
         J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
         J_out = tf.subtract(J_out, J0)
         T_out = tf.subtract(v, J0)