diff --git a/superSMPL/smpl.py b/superSMPL/smpl.py
deleted file mode 100755
index 678f31a..0000000
--- a/superSMPL/smpl.py
+++ /dev/null
@@ -1,350 +0,0 @@
-import os
-import sys
-
-import pickle
-import scipy.io as sio
-from math import *
-
-import numpy as np
-import tensorflow as tf
-
-# Left-Right
-eval_joints = {
-        'All': [1, 2, 3, 4, 5, 7, 8, 9, 12, 15, 16, 17, 18, 19, 20, 21],
-        'Head': [12, 15], # Neck, Head
-        'Torso': [3, 6, 9], # Back joints
-        'Shoulder': [13, 14, 16, 17], # Inner, Outer 
-        'Elbow': [18, 19],
-        'Wrist': [20, 21],
-        'Hip': [1, 2],
-        'Knee': [4, 5],
-        'Foot': [7, 8, 10, 11]} # Ankle, Foot
-        
-class SMPL:
-
-    def __init__(self, path):
-        self.J_regressor = []
-        self.v_template = []
-        self.shapedirs = []
-        self.posedirs = []
-        self.weights = []
-        # Load SMPL models
-        # path = os.path.abspath(os.path.dirname(__file__))
-        with open(os.path.join(path, 'super_model_f.pkl'), 'rb') as f:
-            if sys.version_info[0]<3:
-                female = pickle.load(f)
-            else:
-                female = pickle.load(f, encoding="latin1")
-            self.J_regressor.append(tf.constant(np.array(female['J_regressor'].todense(), dtype=np.float32)))
-            self.v_template.append(tf.constant(female['v_template'], dtype=np.float32))
-            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(path, 'super_model_m.pkl'), 'rb') as f:
-            if sys.version_info[0]<3:
-                male = pickle.load(f)
-            else:
-                male = pickle.load(f, encoding="latin1")
-            self.J_regressor.append(tf.constant(np.array(male['J_regressor'].todense(), dtype=np.float32)))
-            self.v_template.append(tf.constant(male['v_template'], dtype=np.float32))
-            self.posedirs.append(tf.constant(male['posedirs'], dtype=np.float32))
-            self.shapedirs.append(tf.constant(male['shapedirs'], dtype=np.float32))
-            self.weights.append(tf.constant(male['weights'], dtype=np.float32))
-            kintree_table = male['kintree_table']
-            id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])}
-            self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])}
-        self.J_regressor = tf.stack(self.J_regressor, axis=0)
-        self.v_template = tf.stack(self.v_template, axis=0)
-        self.shapedirs = tf.stack(self.shapedirs, axis=0)
-        self.posedirs = tf.stack(self.posedirs, axis=0)
-        self.weights = tf.stack(self.weights, axis=0)
-        self.weights_one_hot = tf.reduce_mean(self.weights, axis=0)
-        self.weights_one_hot = tf.math.argmax(self.weights_one_hot, axis=-1)
-        self.weights_one_hot = tf.one_hot(self.weights_one_hot, 24)
-
-        # Topology no head/hands/feet
-        F = sio.loadmat(path + '/../Data/ssmpl_faces.mat')
-        self.f = np.array(F['faces_default'])
-        self.f_skirt = np.array(F['faces_skirt'])
-        self.nn = np.array(sio.loadmat(path + '/../Data/NN.mat'))
-        self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1)
-
-    def Q2R(self, w, x, y, z):
-        def diag(a, b):  # computes the diagonal entries,  1 - 2*a**2 - 2*b**2
-            return 1 - 2 * tf.pow(a, 2) - 2 * tf.pow(b, 2)
-
-        def tr_add(a, b, c, d):  # computes triangle entries with addition
-            return 2 * a * b + 2 * c * d
-
-        def tr_sub(a, b, c, d):  # computes triangle entries with subtraction
-            return 2 * a * b - 2 * c * d
-
-        m = [[diag(y, z), tr_sub(x, y, z, w), tr_add(x, z, y, w)],
-                [tr_add(x, y, z, w), diag(x, z), tr_sub(y, z, x, w)],
-                [tr_sub(x, z, y, w), tr_add(y, z, x, w), diag(x, y)]]
-            
-        R = tf.stack([tf.stack(m[i], axis=-1) for i in range(3)], axis=-2)
-        return R
-            
-    def rodrigues(self, r):
-        theta = tf.norm(r, axis=(1, 2), keepdims=True)
-        theta = tf.maximum(theta, np.finfo(np.float32).eps)
-        r_hat = r / theta
-        z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float32)
-        m = tf.stack(
-            (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
-            -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), axis=1)
-        m = tf.reshape(m, (-1, 3, 3))
-        i_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros(
-            (theta.get_shape().as_list()[0], 3, 3), dtype=tf.float32)
-        # R = cos * i_cube + (1 - cos) * dot + tf.sin(theta) * m
-        R = i_cube + m * tf.sin(theta) + tf.matmul(m, m) * (1 - tf.cos(theta))
-        return R
-
-    def with_zeros(self, x):
-        return tf.concat((x, tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)), axis=0)
-
-    def pack(self, x):
-        return tf.concat((tf.zeros((x.get_shape().as_list()[0], 4, 3), dtype=tf.float32), x), axis=2)
-        
-    def run_rodrigues(self, pose):
-        batch_size = pose.shape[0]
-        R_out = self.rodrigues(tf.reshape(pose, (-1, 1, 3)))
-        R_out = tf.reshape(R_out, (batch_size, -1, 3, 3))
-        #R_out = []
-        #for i in range(batch_size):
-        #    R_out.append(self.rodrigues(tf.reshape(pose[i,:], (-1, 1, 3))))
-        #R_out = tf.stack(R_out, axis=0)
-        return R_out
-
-    def run(self, R, betas, remove=True):
-        batch_size = betas.shape[0]
-
-        #J_regressor = tf.gather(self.J_regressor, gender, axis=0)
-        #shapedirs = tf.gather(self.shapedirs, gender, axis=0)
-        #v_template = tf.gather(self.v_template, gender, axis=0)
-        
-        v_shaped = []
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
-        v_shaped = tf.stack(v_shaped, axis=-1)
-
-        def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights, parent=self.parent):
-            R = x[0]
-            v_shaped = x[1]
-            g = x[2]
-
-            J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
-            R_cube_big = R
-            R_cube = R_cube_big[1:]
-            I_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((R_cube.get_shape()[0], 3, 3), dtype=tf.float32)
-            lrotmin = tf.squeeze(tf.reshape((R_cube - I_cube), (-1, 1)))
-            v_posed = v_shaped[:, :, g] + tf.tensordot(posedirs[g], lrotmin, axes=[[2], [0]])
-            results = []
-            results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
-            for j in range(1, 24):
-                results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
-            stacked = tf.stack(results, axis=0)
-            results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
-            
-            J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
-            J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
-            
-            T = tf.tensordot(weights[g], results, axes=((1), (0)))
-            rest_shape_h = tf.concat((v_posed, tf.ones((v_posed.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
-            v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
-            v = tf.reshape(v, (-1, 4))[:, :3]
-
-            return v, tf.squeeze(J), v_posed
-
-        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
-        T_out, J_out, v_posed = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, tf.float32))
-        
-        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
-        J_out = tf.subtract(J_out, J0)
-        T_out = tf.subtract(T_out, J0)
-        v_posed = tf.subtract(v_posed, J0)
-        
-        if remove:
-            T_out = tf.gather(T_out, self.to_keep, axis=1)
-            v_posed = tf.gather(v_posed, self.to_keep, axis=1)
-        
-        return T_out, J_out, v_posed
-
-    def run_with_weights(self, R, betas, v_cloth, weights=None, filter=None):
-        
-        v_shaped = []
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
-        v_shaped = tf.stack(v_shaped, axis=-1)
-
-        if weights is None:
-            w = self.weights_one_hot[np.newaxis, :, :]
-            if filter is not None:
-                w = tf.gather(tf.gather(w, self.to_keep, axis=1), filter, axis=1)
-            if v_cloth.shape[1] == 14475:
-                w = tf.gather(w, self.to_keep, axis=1)
-        else:
-            w = weights
-
-        def single_run(x, J_regressor=self.J_regressor, w=w, parent=self.parent):
-            R = x[0]
-            v_shaped = x[1]
-            g = x[2]
-            v_cloth = x[3]
-
-            J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
-            R_cube_big = R
-            results = []
-            results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
-            for j in range(1, 24):
-                results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
-            stacked = tf.stack(results, axis=0)
-            results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
-            
-            J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
-            J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
-            
-            T = tf.tensordot(w[0], results, axes=((1), (0)))
-            rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[0], 1), dtype=tf.float32)), axis=1)
-            v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
-            v = tf.reshape(v, (-1, 4))[:, :3]
-
-            return v, tf.squeeze(J)
-
-        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
-        T_out, J_out = tf.map_fn(single_run, (R, v_shaped, gender, v_cloth), (tf.float32, tf.float32))
-        
-        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
-        T_out = tf.subtract(T_out, J0)
-
-        return T_out
-
-    def get_joints_template(self, betas):
-        v_shaped = []
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
-        v_shaped = tf.stack(v_shaped, axis=-1)
-
-        def single_run(x, J_regressor=self.J_regressor):
-            v_shaped = x[0]
-            g = x[1]
-            J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
-            return J
-
-        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
-        J = tf.map_fn(single_run, (v_shaped, gender), (tf.float32))
-        J0 = tf.reshape(J[:,0,:], (-1, 1, 3))
-        
-        return tf.subtract(J, J0)
-
-
-    def run_garment(self, R, betas, v_cloth, J=None, fl_idx=None):
-        if J is None:
-            J = self.get_joints_template(betas)
-        
-        weights = tf.gather(self.weights, self.to_keep, axis=1)
-        
-        if fl_idx is not None:
-            weights = tf.gather(weights, fl_idx, axis=1)
-
-        def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=weights, parent=self.parent):
-            R = x[0]
-            J = x[1]
-            g = x[2]
-            v_cloth = x[3]
-
-            R_cube_big = R
-            # R_cube = R_cube_big[1:]
-            # I_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((R_cube.get_shape()[0], 3, 3), dtype=tf.float32)
-            # lrotmin = tf.squeeze(tf.reshape((R_cube - I_cube), (-1, 1)))
-            
-            results = []
-            results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
-            for j in range(1, 24):
-                results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
-            stacked = tf.stack(results, axis=0)
-            results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
-            
-            J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
-            J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
-            
-            T = tf.tensordot(weights[g, :, :], results, axes=((1), (0)))
-            rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[0], 1), dtype=tf.float32)), axis=1)
-            v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
-            v = tf.reshape(v, (-1, 4))[:, :3]
-
-            return v, tf.squeeze(J)
-
-        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
-        T_out, J_out = tf.map_fn(single_run, (R, J, gender, v_cloth), (tf.float32, tf.float32))
-        
-        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
-        J_out = tf.subtract(J_out, J0)
-        T_out = tf.subtract(T_out, J0)
-        
-        return T_out
-
-    def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None):
-        
-        if J is None:
-            J = self.get_joints_template(betas)
-        
-        def single_run(x, parent=self.parent):
-            R = x[0]
-            J = x[1]
-            v_cloth = x[2]
-            weights = x[3]
-
-            R_cube_big = R
-       
-            results = []
-            results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
-            for j in range(1, 24):
-                results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
-            stacked = tf.stack(results, axis=0)
-            results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
-            
-            J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
-            J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
-            
-            T = tf.tensordot(weights, results, axes=((1), (0)))
-            rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
-            v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
-            v = tf.reshape(v, (-1, 4))[:, :3]
-
-            return v, tf.squeeze(J)
-
-        # gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
-        T_out, J_out = tf.map_fn(single_run, (R, J, v_cloth, weights), (tf.float32, tf.float32))
-        
-        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
-        J_out = tf.subtract(J_out, J0)
-        T_out = tf.subtract(T_out, J0)
-        
-        return T_out
-
-    def run_backward(self, R, J, v_posed, weights):
-        batch_size = R.get_shape()[0]
-        
-        J_relative = J[:, 1:, :] - tf.gather(J, tf.constant(list(self.parent.values()), tf.int32), axis=1) # batch_size x 23 x 3
-
-        J_out = []
-        J_out.append(tf.expand_dims(J[:, 0, :], axis=-1)) # batch_size x 3 x 1
-        results = []
-        results.append(tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((batch_size, 3, 3), dtype=tf.float32))
-        for j in range(1, 24):
-            results.append(tf.matmul(R[:, self.parent[j], :, :], results[self.parent[j]]))
-            J_out.append(tf.matmul(results[j], tf.expand_dims(J_relative[:, j-1, :], axis=-1)) + J_out[self.parent[j]]) # batch_size x 3 x 1
-        results = tf.transpose(tf.stack(results, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 3)
-        J_out = tf.transpose(tf.stack(J_out, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 1)
-
-        results = tf.matmul(R, results)
-        T = tf.transpose(tf.tensordot(results, weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 3)
-        v_posed = tf.expand_dims(v_posed, axis=-1) # (batch_size x 6890 x 3 x 1)
-        J = tf.expand_dims(J, axis=-1) # (batch_size x 24 x 3 x 1)
-        v = tf.matmul(T, v_posed) + tf.transpose(tf.tensordot(J_out - tf.matmul(results, J), weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 1)
-        v = tf.squeeze(v, -1) # (batch_size x 6890 x 3)
-        J_out = tf.squeeze(J_out, -1) # (batch_size x 24 x 3)
-
-        return v, J_out
\ No newline at end of file
diff --git a/superSMPL/smpl_torch.py b/superSMPL/smpl_torch.py
new file mode 100755
index 0000000..4c0226f
--- /dev/null
+++ b/superSMPL/smpl_torch.py
@@ -0,0 +1,224 @@
+# based on https://github.com/CalciferZh/SMPL/blob/master/smpl_torch_batch.py
+
+import os
+import sys
+import pickle
+import torch
+import numpy as np
+import scipy.io as sio
+from time import time
+from torch.nn import Module
+
+class SMPL(Module):
+  def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False):
+    
+    super(SMPL, self).__init__()
+    self.J_regressor = []
+    self.v_template = []
+    self.shapedirs = []
+    self.posedirs = []
+    self.weights = []
+    self.is_skirt = is_skirt
+    # Load SMPL models
+    # path = os.path.abspath(os.path.dirname(__file__))
+    with open(os.path.join(path, 'super_model_f.pkl'), 'rb') as f:
+        female = pickle.load(f, encoding="latin1")
+        self.J_regressor.append(torch.tensor(np.array(female['J_regressor'].todense(), dtype=np.float32)))
+        self.v_template.append(torch.tensor(female['v_template'], dtype=torch.float32))
+        self.posedirs.append(torch.tensor(female['posedirs'], dtype=torch.float32))
+        self.shapedirs.append(torch.tensor(female['shapedirs'], dtype=torch.float32))
+        self.weights.append(torch.tensor(female['weights'], dtype=torch.float32))
+    with open(os.path.join(path, 'super_model_m.pkl'), 'rb') as f:
+        male = pickle.load(f, encoding="latin1")
+        self.J_regressor.append(torch.tensor(np.array(male['J_regressor'].todense(), dtype=np.float32)))
+        self.v_template.append(torch.tensor(male['v_template'], dtype=torch.float32))
+        self.posedirs.append(torch.tensor(male['posedirs'], dtype=torch.float32))
+        self.shapedirs.append(torch.tensor(male['shapedirs'], dtype=torch.float32))
+        self.weights.append(torch.tensor(male['weights'], dtype=torch.float32))
+        kintree_table = male['kintree_table']
+        id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])}
+        self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])}
+    self.J_regressor = torch.stack(self.J_regressor, axis=0).to(device)
+    self.v_template = torch.stack(self.v_template, axis=0).to(device)
+    self.shapedirs = torch.stack(self.shapedirs, axis=0).to(device)
+    self.posedirs = torch.stack(self.posedirs, axis=0).to(device)
+    self.weights = torch.stack(self.weights, axis=0).to(device)
+
+    # Topology no head/hands/feet
+    F = sio.loadmat(os.path.join(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(os.path.join(path, 'NN.mat')))
+    self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1)
+
+    self.weights_all = self.weights
+    self.weights = self.weights[:,self.to_keep,:]
+    
+    self.shape = (1,27578,3)
+    self.device = device
+
+    if orig_smpl:
+        self.to_keep = self.to_keep[self.to_keep < 6890]
+        self.J_regressor = self.J_regressor[:,:,:6890]
+        self.v_template = self.v_template[:,:6890,:]
+        self.shapedirs = self.shapedirs[:,:6890,:,:]
+        self.posedirs = self.posedirs[:,:6890,:,:]
+        self.weights_all = self.weights_all[:,:6890,:]
+        self.weights = self.weights_all[:,self.to_keep,:]
+        self.shape = (1, 6890, 3)
+
+  @staticmethod
+  def rodrigues(r):
+    """
+    Rodrigues' rotation formula that turns axis-angle tensor into rotation
+    matrix in a batch-ed manner.
+    Parameter:
+    ----------
+    r: Axis-angle rotation tensor of shape [batch_size * angle_num, 1, 3].
+    Return:
+    -------
+    Rotation matrix of shape [batch_size * angle_num, 3, 3].
+    """
+    # eps = r.clone().normal_(std=1e-8)
+    # theta = torch.norm(r + eps, dim=(1, 2), keepdim=True)  # dim cannot be tuple
+    theta = torch.norm(r, dim=(1, 2), keepdim=True)
+    theta = torch.maximum(theta, torch.ones_like(theta) * np.finfo(np.float32).eps)
+
+    theta_dim = theta.shape[0]
+    r_hat = r / theta
+    # cos = torch.cos(theta)
+    z_stick = torch.zeros(theta_dim, dtype=torch.float32).to(r.device)
+    m = torch.stack(
+      (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
+       -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), dim=1)
+    m = torch.reshape(m, (-1, 3, 3))
+    i_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) \
+             + torch.zeros((theta_dim, 3, 3), dtype=torch.float32)).to(r.device)
+    A = r_hat.permute(0, 2, 1)
+    # dot = torch.matmul(A, r_hat)
+    # R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m
+    R = i_cube + m * torch.sin(theta) + torch.matmul(m, m) * (1 - torch.cos(theta))
+    return R
+
+  @staticmethod
+  def with_zeros(x):
+    """
+    Append a [0, 0, 0, 1] tensor to a [3, 4] tensor.
+    Parameter:
+    ---------
+    x: Tensor to be appended.
+    Return:
+    ------
+    Tensor after appending of shape [4,4]
+    """
+    ones = torch.tensor(
+      [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32
+    ).expand(x.shape[0],-1,-1).to(x.device)
+    ret = torch.cat((x, ones), dim=1)
+    return ret
+
+  @staticmethod
+  def pack(x):
+    """
+    Append zero tensors of shape [4, 3] to a batch of [4, 1] shape tensor.
+    Parameter:
+    ----------
+    x: A tensor of shape [batch_size, 4, 1]
+    Return:
+    ------
+    A tensor of shape [batch_size, 4, 4] after appending.
+    """
+    zeros43 = torch.zeros(
+      (x.shape[0], x.shape[1], 4, 3), dtype=torch.float32).to(x.device)
+    ret = torch.cat((zeros43, x), dim=3)
+    return ret
+
+
+  def forward(self, pose, betas, verts: torch.Tensor = None, remove: bool = True):
+    
+    """
+          Construct a compute graph that takes in parameters and outputs a tensor as
+          model vertices. Face indices are also returned as a numpy ndarray.
+          
+          20190128: Add batch support.
+          Parameters:
+          ---------
+          pose: Also known as 'theta', an [N, 24, 3] tensor indicating child joint rotation
+          relative to parent joint. For root joint it's global orientation.
+          Represented in a axis-angle format.
+          betas: Parameter for model shape. A tensor of shape [N, 10] as coefficients of
+          PCA components. Only 10 components were released by SMPL author.
+          trans: Global translation tensor of shape [N, 3].
+          Return:
+          ------
+          A 3-D tensor of [N * 6890 * 3] for vertices,
+          and the corresponding [N * 19 * 3] joint positions.
+    """
+    batch_num = betas.shape[0]
+
+    # v_shaped = torch.tensordot(betas, self.shapedirs, dims=([1], [2])) + self.v_template
+    v_shaped = []
+    v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[0], dims=[[1], [2]]) + torch.reshape(self.v_template[0], shape=self.shape)) # female
+    v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[1], dims=[[1], [2]]) + torch.reshape(self.v_template[1], shape=self.shape)) # male
+    v_shaped = torch.stack(v_shaped, axis=-1)
+    
+    # J = torch.matmul(self.J_regressor, v_shaped)
+    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])])
+    R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
+
+    if verts is not None:
+      v_posed = verts
+    else:
+      R_cube = R_cube_big[:, 1:, :, :]
+      I_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) + \
+        torch.zeros((batch_num, R_cube.shape[1], 3, 3), dtype=torch.float32)).to(self.device)
+      lrotmin = (R_cube - I_cube).reshape(batch_num, -1, 1).squeeze(dim=2)
+      v_posed = torch.stack([v_shaped[i, :, :, g.to(torch.long)] + torch.tensordot(lrotmin[i], self.posedirs[g.to(torch.long)], dims=([0], [2])) for i, g in enumerate(betas[:,-1])])
+
+    results = []
+    results.append(
+      self.with_zeros(torch.cat((R_cube_big[:, 0], torch.reshape(J[:, 0, :], (-1, 3, 1))), dim=2))
+    )
+    for i in range(1, 24):
+      results.append(
+        torch.matmul(
+          results[self.parent[i]],
+          self.with_zeros(
+            torch.cat(
+              (R_cube_big[:, i], torch.reshape(J[:, i, :] - J[:, self.parent[i], :], (-1, 3, 1))),
+              dim=2
+            )
+          )
+        )
+      )
+    
+    stacked = torch.stack(results, dim=1)
+    results = stacked - \
+      self.pack(
+        torch.matmul(
+          stacked,
+          torch.reshape(
+            torch.cat((J, torch.zeros((batch_num, 24, 1), dtype=torch.float32).to(self.device)), dim=2),
+            (batch_num, 24, 4, 1)
+          )
+        )
+      )
+    
+    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[g.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
+    )
+    v = torch.matmul(T, torch.reshape(rest_shape_h, (batch_num, -1, 4, 1)))
+    v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
+    result = v
+
+    root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
+    result = result - root_joint
+    joints = joints - root_joint
+
+    if remove and verts is None:
+      result = result[:,self.to_keep,:]
+
+    return result, joints