hot fixes so it can run in the cluster

src/diffwave/__main__.py

@@ -21,7 +21,6 @@ from torch.multiprocessing import spawn
 import os
 import sys
 SOURCE_DIR = os.path.abspath(os.path.join(os.path.dirname( __file__ ), '../..', 'src'))
-print(SOURCE_DIR)
 sys.path.append(SOURCE_DIR)
 
 from learner import train, train_distributed

src/diffwave/preprocess.py

@@ -23,6 +23,12 @@ from concurrent.futures import ProcessPoolExecutor
 from glob import glob
 from tqdm import tqdm
 
+# Needed to import from src/ in general
+import os
+import sys
+SOURCE_DIR = os.path.abspath(os.path.join(os.path.dirname( __file__ ), '../..', 'src'))
+sys.path.append(SOURCE_DIR)
+
 from diffwave.params import params