updated dockerfile and training.md with correct jax-tpu-emb file

Author: ajkv-google

Dockerfile

@@ -25,8 +25,8 @@ RUN pip install -r ./requirements.txt
 RUN pip install "protobuf>=6.31.1"
 
 # Install the latest jax-tpu-embedding wheel
-COPY jax_tpu_embedding-0.1.0.dev20260121-cp312-cp312-manylinux_2_31_x86_64.whl ./
-RUN pip install ./jax_tpu_embedding-0.1.0.dev20260121-cp312-cp312-manylinux_2_31_x86_64.whl
+COPY jax_tpu_embedding-0.1.0.dev20260226-cp312-cp312-manylinux_2_31_x86_64.whl ./
+RUN pip install ./jax_tpu_embedding-0.1.0.dev20260226-cp312-cp312-manylinux_2_31_x86_64.whl
 
 # Copy the current directory contents into the container
 COPY . /app

training.md

@@ -26,7 +26,7 @@ source venv/bin/activate
 
 Install the latest version of the jax-tpu-embedding library:
 ```bash
-pip install ./jax_tpu_embedding-0.1.0.dev20260121-cp312-cp312-manylinux_2_31_x86_64.whl
+pip install ./jax_tpu_embedding-0.1.0.dev20260226-cp312-cp312-manylinux_2_31_x86_64.whl
 ```
 ```bash
 pip install -r requirements.txt