updates from main

.gitignore

@@ -7,4 +7,5 @@ temp/
 *.hash
 *.npz
 sweep_results/
-.gradio
+.gradio
+__pycache__

Dockerfile

@@ -1,6 +1,6 @@
 FROM zeahub/all:v0.0.4
 
-RUN pip install --no-cache-dir tyro optuna
+RUN pip install --no-cache-dir gradio tyro optuna
 
 RUN pip install --no-cache-dir --no-deps pytorch_fid
 

main.py

@@ -1,6 +1,9 @@
 import copy
+import os
 from pathlib import Path
 
+os.environ["KERAS_BACKEND"] = "jax"
+
 import jax
 import keras
 import matplotlib.pyplot as plt