diff --git a/models.py b/models.py
index 97543be..c9e1467 100644
--- a/models.py
+++ b/models.py
@@ -190,13 +190,14 @@ class DeepHealth(nn.Module):
         daily_mask = flatten_mask(exposure_daily_mask, exposure_daily)
         monthly_mask = flatten_mask(exposure_monthly_mask, exposure_monthly)
 
-        param = next(self.exposure_encoder.parameters())
-        daily = daily.to(device=param.device, dtype=param.dtype)
-        monthly = monthly.to(device=param.device, dtype=param.dtype)
+        exposure_device = exposure_daily.device
+        exposure_dtype = self.token_embedding.weight.dtype
+        daily = daily.to(device=exposure_device, dtype=exposure_dtype)
+        monthly = monthly.to(device=exposure_device, dtype=exposure_dtype)
         if daily_mask is not None:
-            daily_mask = daily_mask.to(device=param.device)
+            daily_mask = daily_mask.to(device=exposure_device)
         if monthly_mask is not None:
-            monthly_mask = monthly_mask.to(device=param.device)
+            monthly_mask = monthly_mask.to(device=exposure_device)
 
         exposure_emb = self.exposure_encoder(
             daily=daily,