JiaruiLi/DeepHealthExpo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix exposure encoding in DataParallel

Browse Source
This commit is contained in:
Jiarui Li
2026-07-09 10:47:23 +08:00
parent 4c60dbb9d9
commit 8976f1ed89
1 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
models.py
View File

@@ -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,