From 9fabdf3d0636188b4a1e38319cecef6ef387420c Mon Sep 17 00:00:00 2001 From: d3sm0 Date: Mon, 17 Jan 2022 09:25:23 -0500 Subject: [PATCH] Ensure compatibility #109 --- experiment_buddy/experiment_buddy.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/experiment_buddy/experiment_buddy.py b/experiment_buddy/experiment_buddy.py index 57b9795..363752c 100644 --- a/experiment_buddy/experiment_buddy.py +++ b/experiment_buddy/experiment_buddy.py @@ -2,6 +2,7 @@ import datetime import logging import os +import re import subprocess import sys import time @@ -257,7 +258,8 @@ def ensure_torch_compatibility(): with open("requirements.txt") as fin: reqs = fin.read() # torch, vision or audio. - if "torch" not in reqs and "torch==1.7.1+cu110" not in reqs: + matches = re.search(r"torch==.*cu.*", reqs) + if "torch" in reqs and not len(matches): # https://mila-umontreal.slack.com/archives/CFAS8455H/p1624292393273100?thread_ts=1624290747.269100&cid=CFAS8455H warnings.warn("""torch rocm4.2 version will be installed on the cluster which is not supported specify torch==1.7.1+cu110 in requirements.txt instead""")