From 4a549379478f0a12823bf772af25e71628ff1629 Mon Sep 17 00:00:00 2001 From: rainx0r Date: Thu, 22 Aug 2024 13:41:59 +0100 Subject: [PATCH] Revert image to custom --- metaworld-jax/Dockerfile | 4 ++-- metaworld-jax/requirements.txt | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/metaworld-jax/Dockerfile b/metaworld-jax/Dockerfile index d0507e0..9c95316 100644 --- a/metaworld-jax/Dockerfile +++ b/metaworld-jax/Dockerfile @@ -1,7 +1,7 @@ -FROM ghcr.io/nvidia/jax:jax-2024-08-19 +FROM python:3.12.5-slim LABEL maintainer="me@evangelos.ai" -RUN apt-get update && apt install -y --no-install-recommends git libglfw3 libglfw3-dev +RUN apt-get update && apt install -y --no-install-recommends git python3-pip libglfw3 libglfw3-dev WORKDIR /usr/src/app COPY requirements.txt ./ diff --git a/metaworld-jax/requirements.txt b/metaworld-jax/requirements.txt index f84e4c3..c319a77 100644 --- a/metaworld-jax/requirements.txt +++ b/metaworld-jax/requirements.txt @@ -1,11 +1,13 @@ # Jax +jax[cuda12]==0.4.31 +flax==0.8.5 distrax==0.1.5 # Metaworld metaworld @ git+https://github.com/rainx0r/Metaworld.git@f131964 # Pytorch CPU for logging ---extra-index-url https://download.pytorch.org/whl/cpu +# --extra-index-url https://download.pytorch.org/whl/cpu torch==2.4.0 # Logging