cast background to device (#2381)

* cast background_color to device

* make onliner instead
This commit is contained in:
Maxim Bonnaerens authored and GitHub committed 2023-08-29 15:07:43 -07:00
1 parent d56338bce1
commit 242c23f0f0
1 file changed
+1 -1
+1 -1
View File
@@ -116,7 +116,7 @@ class RGBRenderer(nn.Module):
background_color = rgb[..., -1, :]
assert isinstance(background_color, torch.Tensor)
comp_rgb = comp_rgb + background_color * (1.0 - accumulated_weight)
comp_rgb = comp_rgb + background_color.to(comp_rgb.device) * (1.0 - accumulated_weight)
return comp_rgb
@classmethod