From 29018c7e2a3b5a9c16b04f4874f27b2334a2b877 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20B=C3=B6ther?=
 <2116466+MaxiBoether@users.noreply.github.com>
Date: Wed, 19 Jun 2024 11:50:03 +0200
Subject: [PATCH] Fix device for initial RHO-LOSS tensor (#525)

The initial tensor resides on CPU. For me, it failed when running
RHO-LOSS, since the very first ` self.rho_loss =
torch.cat([self.rho_loss, training_loss -
irreducible_loss]).to(training_loss.dtype)` was on two different
devices, the CPU (initial tensor) and the GPU (training loss and IR
loss). This PR fixes that by moving the initial tensor to the correct
device.
---
 .../remote_downsamplers/remote_rho_loss_downsampling.py        | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/modyn/trainer_server/internal/trainer/remote_downsamplers/remote_rho_loss_downsampling.py b/modyn/trainer_server/internal/trainer/remote_downsamplers/remote_rho_loss_downsampling.py
index 4b707f5a4..c1ccb84f2 100644
--- a/modyn/trainer_server/internal/trainer/remote_downsamplers/remote_rho_loss_downsampling.py
+++ b/modyn/trainer_server/internal/trainer/remote_downsamplers/remote_rho_loss_downsampling.py
@@ -36,10 +36,11 @@ def __init__(
         )
         self.rho_loss: torch.Tensor = torch.tensor([])
         self.number_of_points_seen = 0
+        self._device = device
 
     def init_downsampler(self) -> None:
         self.index_sampleid_map: list[int] = []
-        self.rho_loss = torch.tensor([])
+        self.rho_loss = torch.tensor([]).to(self._device)
         self.number_of_points_seen = 0
 
     def inform_samples(