From 5330e1eb7728c6e708c6e7f445a0876037c69c16 Mon Sep 17 00:00:00 2001
From: Wesley Gifford <79663411+wgifford@users.noreply.github.com>
Date: Tue, 10 Dec 2024 13:46:47 -0500
Subject: [PATCH 1/2] explicitly set device

---
 services/inference/tsfminference/hf_service_handler.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/services/inference/tsfminference/hf_service_handler.py b/services/inference/tsfminference/hf_service_handler.py
index 1cab7417..3eb7459b 100644
--- a/services/inference/tsfminference/hf_service_handler.py
+++ b/services/inference/tsfminference/hf_service_handler.py
@@ -9,6 +9,7 @@
 from typing import Any, Dict, Optional, Union
 
 import pandas as pd
+import torch
 import transformers
 from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
 
@@ -353,12 +354,14 @@ def _run(
                     total_periods=model_prediction_length,
                 )
 
+        device = "cpu" if not torch.cuda.is_available() else "cuda"
         forecast_pipeline = TimeSeriesForecastingPipeline(
             model=self.model,
             explode_forecasts=True,
             feature_extractor=self.preprocessor,
             add_known_ground_truth=False,
             freq=self.preprocessor.freq,
+            device=device,
         )
         forecasts = forecast_pipeline(data, future_time_series=future_data, inverse_scale_outputs=True)
 

From c96905dba5cc3c1943c5360cecf70400ced5a139 Mon Sep 17 00:00:00 2001
From: Wesley Gifford <79663411+wgifford@users.noreply.github.com>
Date: Tue, 10 Dec 2024 14:04:18 -0500
Subject: [PATCH 2/2] select device

---
 services/finetuning/tsfmfinetuning/finetuning.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/services/finetuning/tsfmfinetuning/finetuning.py b/services/finetuning/tsfmfinetuning/finetuning.py
index 374048a2..4b63e7d5 100644
--- a/services/finetuning/tsfmfinetuning/finetuning.py
+++ b/services/finetuning/tsfmfinetuning/finetuning.py
@@ -7,6 +7,7 @@
 from typing import Any, Dict, Tuple, Union
 
 import pandas as pd
+import torch
 from fastapi import APIRouter, HTTPException
 from starlette import status
 from transformers import EarlyStoppingCallback, Trainer, TrainingArguments, set_seed
@@ -215,7 +216,7 @@ def _finetuning_common(
             metric_for_best_model="eval_loss",  # Metric to monitor for early stopping
             greater_is_better=False,  # For loss
             label_names=["future_values"],
-            use_cpu=True,  # only needed for testing on Mac :(
+            use_cpu=not torch.cuda.is_available(),
         )
 
         callbacks = []