From baf444b8615bf3a90f31c8ff0e95f46450ef8ff8 Mon Sep 17 00:00:00 2001
From: Tim O'Donnell <timodonnell@gmail.com>
Date: Tue, 18 Jun 2019 11:14:40 -0400
Subject: [PATCH] fix
---
mhcflurry/train_pan_allele_models_command.py | 2 --
test/test_train_and_related_commands.py | 2 ++
test/test_train_pan_allele_models_command.py | 2 ++
3 files changed, 4 insertions(+), 2 deletions(-)
diff --git a/mhcflurry/train_pan_allele_models_command.py b/mhcflurry/train_pan_allele_models_command.py
index ca488076..56e4d60b 100644
--- a/mhcflurry/train_pan_allele_models_command.py
+++ b/mhcflurry/train_pan_allele_models_command.py
@@ -32,8 +32,6 @@ from .allele_encoding import AlleleEncoding
from .encodable_sequences import EncodableSequences
from .regression_target import to_ic50, from_ic50
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
-
# To avoid pickling large matrices to send to child processes when running in
# parallel, we use this global variable as a place to store data. Data that is
diff --git a/test/test_train_and_related_commands.py b/test/test_train_and_related_commands.py
index 197ade0d..f6abfd26 100644
--- a/test/test_train_and_related_commands.py
+++ b/test/test_train_and_related_commands.py
@@ -14,6 +14,8 @@ from numpy.testing import assert_array_less, assert_equal
from mhcflurry import Class1AffinityPredictor
from mhcflurry.downloads import get_path
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
HYPERPARAMETERS = [
{
"n_models": 2,
diff --git a/test/test_train_pan_allele_models_command.py b/test/test_train_pan_allele_models_command.py
index 870c3694..25a9bdc7 100644
--- a/test/test_train_pan_allele_models_command.py
+++ b/test/test_train_pan_allele_models_command.py
@@ -18,6 +18,8 @@ from mhcflurry import Class1AffinityPredictor,Class1NeuralNetwork
from mhcflurry.allele_encoding import AlleleEncoding
from mhcflurry.downloads import get_path
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
HYPERPARAMETERS_LIST = [
{
--
GitLab