diff --git a/build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/amd.sh b/build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/amd.sh
index 7824822..b86c74b 100755
--- a/build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/amd.sh
+++ b/build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/amd.sh
@@ -10,7 +10,7 @@ build_amd_install_torch() {
         torch==${PYTORCH_VERSION} \
         torchvision \
         torchaudio \
-        --extra-index-url=https://download.pytorch.org/whl/rocm${ROCM_VERSION}
+        --index-url=https://download.pytorch.org/whl/rocm${ROCM_VERSION}
 }
 
 build_amd_main "$@"
\ No newline at end of file
diff --git a/build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/nvidia.sh b/build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/nvidia.sh
index b7c08bd..02996f0 100755
--- a/build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/nvidia.sh
+++ b/build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/nvidia.sh
@@ -9,12 +9,12 @@ build_nvidia_main() {
 build_nvidia_install_torch() {
     short_cuda_version="cu$(cut -d '.' -f 1,2 <<< "${CUDA_VERSION}" | tr -d '.')"
     "$KOHYA_VENV_PIP" install --no-cache-dir \
-        nvidia-ml-py3 \
         torch==${PYTORCH_VERSION} \
         torchvision \
         torchaudio \
         xformers \
-        --extra-index-url=https://download.pytorch.org/whl/$short_cuda_version
+        --index-url=https://download.pytorch.org/whl/$short_cuda_version
+    
 }
 
 build_nvidia_run_tests() {
diff --git a/build/COPY_ROOT_1/opt/ai-dock/bin/preflight.d/10-default.sh b/build/COPY_ROOT_1/opt/ai-dock/bin/preflight.d/10-default.sh
index 7ff7b09..468d917 100755
--- a/build/COPY_ROOT_1/opt/ai-dock/bin/preflight.d/10-default.sh
+++ b/build/COPY_ROOT_1/opt/ai-dock/bin/preflight.d/10-default.sh
@@ -6,6 +6,9 @@ function preflight_main() {
     preflight_configure_accelerate
     preflight_update_kohya_ss
     printf "%s" "${KOHYA_FLAGS}" > /etc/kohya_ss_flags.conf
+    export TENSORBOARD_FLAGS=${TENSORBOARD_FLAGS:-"--logdir /opt/kohya_ss/logs"}
+    env-store TENSORBOARD_FLAGS
+    printf "%s" "${TENSORBOARD_FLAGS}" > /etc/tensorboard_flags.conf
 }
 
 function preflight_configure_accelerate() {
diff --git a/build/COPY_ROOT_1/opt/ai-dock/bin/supervisor-tensorboard.sh b/build/COPY_ROOT_1/opt/ai-dock/bin/supervisor-tensorboard.sh
index 42956fe..c05f2d9 100755
--- a/build/COPY_ROOT_1/opt/ai-dock/bin/supervisor-tensorboard.sh
+++ b/build/COPY_ROOT_1/opt/ai-dock/bin/supervisor-tensorboard.sh
@@ -66,9 +66,10 @@ function start() {
     printf "Starting %s...\n" "${SERVICE_NAME}"
     
     source "$KOHYA_VENV/bin/activate"
+    FLAGS_COMBINED="$(cat /etc/tensorboard_flags.conf)"
     LD_PRELOAD=libtcmalloc.so tensorboard \
-        --port ${LISTEN_PORT} \
-        --logdir /opt/kohya_ss/logs
+        $FLAGS_COMBINED \
+        --port ${LISTEN_PORT}
 }
 
 start 2>&1
\ No newline at end of file
diff --git a/docker-compose.yaml b/docker-compose.yaml
index f4df640..13a4fb0 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -74,6 +74,7 @@ services:
         - SYNCTHING_UI_PORT_HOST=${SYNCTHING_UI_PORT_HOST:-8384}
         - SYNCTHING_TRANSPORT_PORT_HOST=${SYNCTHING_TRANSPORT_PORT_HOST:-22999}
         - SYNCTHING_URL=${SYNCTHING_URL:-}
+        - TENSORBOARD_FLAGS=${TENSORBOARD_FLAGS:-"--logdir /opt/kohya_ss/logs"}
         - TENSORBOARD_PORT_HOST=${TENSORBOARD_PORT_HOST:-6006}
         - TENSORBOARD_PORT_LOCAL=${TENSORBOARD_PORT_LOCAL:-16006}
         - TENSORBOARD_METRICS_PORT=${TENSORBOARD_METRICS_PORT:-26006}