| version: '3.8' | |
| services: | |
| toxic-classifier: | |
| build: . | |
| runtime: nvidia # Enable NVIDIA runtime for GPU support | |
| environment: | |
| - NVIDIA_VISIBLE_DEVICES=all | |
| - WANDB_API_KEY=${WANDB_API_KEY} # Set this in .env file | |
| volumes: | |
| - ./dataset:/app/dataset # Mount dataset directory | |
| - ./weights:/app/weights # Mount weights directory | |
| command: python model/train.py # Default command, can be overridden |