From 8976f1ed89d54fb9b20fc1555ef37978a5c3eea6 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Thu, 9 Jul 2026 10:47:23 +0800 Subject: [PATCH] Fix exposure encoding in DataParallel --- models.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) 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,