From 4f8c594daae189a834c09fdec4a2f620bbf59c64 Mon Sep 17 00:00:00 2001 From: sileod Date: Fri, 20 Jan 2023 09:13:46 +0100 Subject: [PATCH] allow cpu --- src/tasknet/models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/tasknet/models.py b/src/tasknet/models.py index c72edd9..9a120a9 100755 --- a/src/tasknet/models.py +++ b/src/tasknet/models.py @@ -63,7 +63,6 @@ def __init__(self, tasks, args, warm_start=None): self.shared_encoder = warm_start self.models={} self.task_names = [t.name for t in tasks] - task_models_list = [] for i, task in progress(list(enumerate(tasks))): model_type = eval(f"AutoModelFor{task.task_type}") @@ -98,9 +97,10 @@ def __init__(self, tasks, args, warm_start=None): self.task_models_list = nn.ModuleList(task_models_list) + device = torch.cuda.current_device() if torch.cuda.is_available() else "cpu" self.Z = nn.parameter.Parameter( torch.zeros(len(tasks), - self.shared_encoder.config.hidden_size, device=torch.cuda.current_device()), + self.shared_encoder.config.hidden_size, device=device), requires_grad=len(tasks)>1 )