Update prediction.py

prediction.py

@@ -6,6 +6,9 @@ import torch.nn.functional as F
 from torchvision import transforms
 from pathlib import Path
 from common import read_yaml
+import os
+os.environ['MPLCONFIGDIR'] = '/tmp'
+
 
 PARAMS_FILE_PATH = Path("params.yaml")