version: '3.8'

services:
  pytorch-training:
    build:
      context: ..
      dockerfile: docker/Dockerfile.pytorch
    volumes:
      - ../:/app
      - ../data:/app/data
      - ../models:/app/models
    environment:
      - PYTHONUNBUFFERED=1
      - CUDA_VISIBLE_DEVICES=0
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    shm_size: '8gb'