diff --git a/docker-compose.yml b/docker-compose.yml
index 6eae9e4..acc0ee4 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -5,7 +5,7 @@ services:
     network_mode: bridge
     environment:
       - HF_TOKEN=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
-      - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+      - PYTORCH_ALLOC_CONF=expandable_segments:True,throw_on_cudamalloc_oom:True
     ipc: host
     ports:
       - "8000:8000"