diff --git a/Dockerfile b/Dockerfile
index 57db175..cd5d9f3 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,5 +1,5 @@
 # Use an existing image as the base
-FROM pytorch/pytorch:1.12.0-cuda11.3-cudnn8-devel
+FROM nvcr.io/nvidia/pytorch:23.03-py3
 
 # Add the current user to the image
 ARG USER_NAME=lhunor
@@ -8,7 +8,7 @@ ARG GROUP_NAME=guest
 ARG GROUP_ID=11000
 
 # Update the system and install required packages
-RUN apt-get update && apt-get upgrade -y && apt-get install -y sudo git
+RUN apt-get update && apt-get upgrade -y && apt-get install -y sudo git apt-utils
 
 # Create a new group with the same GID as the host
 RUN groupadd --gid $GROUP_ID $GROUP_NAME
@@ -26,8 +26,10 @@ RUN usermod -aG sudo $USER_NAME
 # Set the default user to the new user
 USER $USER_NAME
 
-WORKDIR /home/$USER_NAME
+WORKDIR /workspace
 
+ENV PATH="$PATH:/home/$USER_NAME/.local/bin"
+RUN pip install --upgrade pip
 COPY requirements.txt /tmp/
 RUN pip install --requirement /tmp/requirements.txt
 
diff --git a/docker-compose.yml b/docker-compose.yml
new file mode 100644
index 0000000..e12e301
--- /dev/null
+++ b/docker-compose.yml
@@ -0,0 +1,21 @@
+services:
+  lhunor_torch:
+    image: lhunor_smpl:v2
+    container_name: lhunor_smpl
+    #command: python -c "import tensorflow as tf;tf.test.gpu_device_name()"
+    ipc: host
+    tty: true
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    volumes:
+      - '~/workspace:/workspace'
+      - '/data-net/guest/lhunor/data:/data-net'
+      - '/data-fast/126-data1/lhunor/data:/data-local/126'
+      - '/data-fast/127-data1/lhunor/data:/data-local/127'
+      - '/data-fast/128-data1/lhunor/data:/data-local/128'
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]
\ No newline at end of file
diff --git a/smpl_tf.py b/smpl_tf.py
index 9091a39..9503b82 100644
--- a/smpl_tf.py
+++ b/smpl_tf.py
@@ -141,6 +141,7 @@ class SMPL:
     def pack(self, x):
         return tf.concat((tf.zeros((x.get_shape().as_list()[0], 4, 3), dtype=tf.float32), x), axis=2)
         
+    # @tf.function
     def run_rodrigues(self, pose):
         batch_size = pose.shape[0]
         R_out = self.rodrigues(tf.reshape(pose, (-1, 1, 3)))
@@ -151,6 +152,7 @@ class SMPL:
         #R_out = tf.stack(R_out, axis=0)
         return R_out
 
+    # @tf.function
     def run(self, R, betas, remove=True):
         batch_size = betas.shape[0]
         # R = self.run_rodrigues(pose)
@@ -182,14 +184,17 @@ class SMPL:
             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)
+            # 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]
 
+            J = tf.transpose(tf.tensordot(v, self.J_regressor[g], axes=([0], [1])), [1,0])
+
             return v, tf.squeeze(J), v_posed, lrotmin
 
         gender = tf.cast(tf.greater_equal(betas[:, 10] , 0.5), tf.int32)
diff --git a/test.py b/test.py
index 9657e2f..90d04f0 100644
--- a/test.py
+++ b/test.py
@@ -1,10 +1,10 @@
 import os
 
-os.environ["CUDA_VISIBLE_DEVICES"] = '5'
+os.environ["CUDA_VISIBLE_DEVICES"] = '7'
 
 import numpy as np
-from smpl_tf import SMPL as SMPL_tf
 from smpl_torch import SMPL as SMPL_torch
+from smpl_tf import SMPL as SMPL_tf
 from utils import plot_mesh, plot_mesh_2
 from time import time
 
@@ -17,7 +17,7 @@ if __name__ == '__main__':
     smpl_tf : SMPL_tf = SMPL_tf('model_files')
     smpl_torch : SMPL_torch = SMPL_torch(device=device, path='model_files')
 
-
+    np.random.seed(4444)
     pose = ((np.random.rand(4, 72) - 0.5) * 0.4).astype(np.float32)
     betas = ((np.random.rand(4, 10) - 0.5) * 0.06).astype(np.float32)
     betas = np.concatenate((betas, np.array([0,1,0,1], dtype=np.float32)[:,None]), axis=1)
@@ -37,10 +37,10 @@ if __name__ == '__main__':
     # plot_mesh(result[0].cpu().numpy(), model.f).show()
 
     R = smpl_tf.run_rodrigues(pose)
-    result_tf, joints_torch, _, _ = smpl_tf.run(R, betas)
+    result_tf, joints_tf, _, _ = smpl_tf.run(R, betas)
     s = time()
     R = smpl_tf.run_rodrigues(pose)
-    result_tf, joints_torch, v_shaped_2, v_posed_2 = smpl_tf.run(R, betas)
+    result_tf, joints_tf, v_shaped_2, v_posed_2 = smpl_tf.run(R, betas)
     print(time() - s)
 
     result_torch = result_torch.cpu().numpy()