version: '2.0'
services:
  main:
    image: tverous/pytorch-notebook
    environment:
      - JUPYTER_TOKEN=passwd # TODO: Change the token
    volumes:
      - /local_vol:/docker_vol
    ports:
      - 8888:8888
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]