Fix: runtime error, workload was not healthy after 30 min

Dockerfile

@@ -48,9 +48,7 @@ RUN mkdir -p /tmp/matplotlib /tmp/numba_cache && \
 # Switch to the non-root user
 USER user
 
-# Configure Accelerate for the user (default to fp16 for speed)
-# This writes to ~/.cache/huggingface/accelerate/default_config.yaml
-RUN python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')"
+# Accelerate configuration is now handled in entrypoint.sh at runtime
 
 ENTRYPOINT ["/app/entrypoint.sh"]
 CMD ["--energy_head_enabled", "--loss_type", "energy_contrastive", "--push_to_hub", "--hub_model_id", "Uday/ctm-energy-based-halting"]

README.md

@@ -224,12 +224,25 @@ If you use this code or build upon CTM in your work, please cite:
 ```bibtex
 @article{ctm2025,
   title={The Continuous Thought Machine},
-  author={...},
+  author={Darlow, Luke and Regan, Ciaran and Risi, Sebastian and Seely, Jeffrey and Jones, Llion},
   journal={arXiv preprint arXiv:2505.05522},
   year={2025}
 }
 ```
 
+### Energy-Based Halting Extension
+
+This repository contains experimental extensions for Energy-Based Halting developed by **Uday Phalak**.
+
+```bibtex
+@misc{ctmenergy2025,
+  title={Energy-Based Halting for Continuous Thought Machines},
+  author={Phalak, Uday},
+  year={2025},
+  note={Experimental Extension}
+}
+```
+
 ---
 
 ## 📝 License

entrypoint.sh

@@ -15,5 +15,9 @@ if [ -n "$HF_TOKEN" ]; then
     args+=("--hub_token" "$HF_TOKEN")
 fi
 
+# Generate Accelerate config at runtime to detect GPUs correctly
+# This writes to ~/.cache/huggingface/accelerate/default_config.yaml
+python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')"
+
 # Run accelerate launch with the training script and arguments
 exec accelerate launch tasks/image_classification/train_energy.py "${args[@]}"