Add GPU support

README.md

@@ -1,7 +1,17 @@
 
-wget https://zenodo.org/record/7044411/files/Clean.zip?download=1 Clean.zip
+## Install Packages
+`python3 -m venv env`
+`pip install -e .`
 
-unzip Clean.zip
+## Download [GuitarFX Dataset] (https://zenodo.org/record/7044411/)
+`./download_egfx.sh`
+
+## Train model
+1. Change Wandb variables in `shell_vars.sh`
+2. `python train.py exp=audio_diffusion`
+or
+2. `python train.py exp=umx`
+
+To add gpu, add `trainer.accelerator='gpu' trainer.devices=-1` to the command-line
+Ex. `python train.py exp=umx trainer.accelerator='gpu' trainer.devices=-1`
 
-python3 -m venv env
-pip install -e .

config.yaml

@@ -48,3 +48,5 @@ trainer:
   enable_model_summary: False
   log_every_n_steps: 1 # Logs metrics every N batches
   accumulate_grad_batches: 1
+  accelerator: null
+  devices: 1