Remove redundant calls to change device

src/model.py

@@ -1,6 +1,6 @@
 from huggingface_hub import hf_hub_download
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from shared import CustomTokens, device
+from shared import CustomTokens
 from errors import ClassifierLoadError, ModelLoadError
 from functools import lru_cache
 import pickle
@@ -100,7 +100,6 @@ def get_model_tokenizer(model_name_or_path, cache_dir=None):
     # Load pretrained model and tokenizer
     model = AutoModelForSeq2SeqLM.from_pretrained(
         model_name_or_path, cache_dir=cache_dir)
-    model.to(device())
 
     tokenizer = AutoTokenizer.from_pretrained(
         model_name_or_path, max_length=model.config.d_model, cache_dir=cache_dir)

src/predict.py

@@ -10,7 +10,7 @@ import logging
 import os
 import itertools
 from utils import re_findall
-from shared import CustomTokens, START_SEGMENT_TEMPLATE, END_SEGMENT_TEMPLATE, OutputArguments, device, seconds_to_time
+from shared import CustomTokens, START_SEGMENT_TEMPLATE, END_SEGMENT_TEMPLATE, OutputArguments, seconds_to_time
 from typing import Optional
 from segment import (
     generate_segments,
@@ -301,7 +301,7 @@ CATEGORIES = [None, 'SPONSOR', 'SELFPROMO', 'INTERACTION']
 def predict_sponsor_text(text, model, tokenizer):
     """Given a body of text, predict the words which are part of the sponsor"""
     input_ids = tokenizer(
-        f'{CustomTokens.EXTRACT_SEGMENTS_PREFIX.value} {text}', return_tensors='pt', truncation=True).input_ids.to(device())
+        f'{CustomTokens.EXTRACT_SEGMENTS_PREFIX.value} {text}', return_tensors='pt', truncation=True).input_ids
 
     max_out_len = round(min(
         max(

src/shared.py

@@ -107,10 +107,6 @@ class GeneralArguments:
         torch.cuda.manual_seed_all(self.seed)
 
 
-def device():
-    return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-
 def seconds_to_time(seconds, remove_leading_zeroes=False):
     fractional = round(seconds % 1, 3)
     fractional = '' if fractional == 0 else str(fractional)[1:]