FROM tensorflow/tensorflow:latest-gpu

# (Optional) Install JupyterLab
RUN pip install jupyterlab \
    && apt update \
    && apt install -y --no-install-recommends nodejs \
    && apt clean \
    && rm -rf /var/lib/apt/lists/*

# # (Optional) Add virtual env in Jupyter
# RUN source <venv_bin_path>/activate \
#     && pip install ipykernel \
#     && python -m ipykernel install --user --name <your_venv_name> \
#     && deactivate

WORKDIR /workspace