GPU bug fixed

Signed-off-by: Devansh Agarwal <[email protected]>

Author: code-dev05

src/anomalib/models/image/glass/lightning_model.py

@@ -170,8 +170,6 @@ def __init__(
         self.c = torch.tensor([1])
         self.lr = lr
 
-        self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
         if pre_proj > 0:
             self.proj_opt = optim.AdamW(
                 self.model.pre_projection.parameters(),
@@ -275,7 +273,7 @@ def training_step(
             self.backbone_opt.zero_grad()
 
         img = batch.image
-        true_loss, gaus_loss, bce_loss, focal_loss, loss = self.model(img, self.c)
+        true_loss, gaus_loss, bce_loss, focal_loss, loss = self.model(img, self.c, self.device)
         loss.backward()
 
         if self.proj_opt is not None:
@@ -301,9 +299,9 @@ def on_train_start(self) -> None:
         with torch.no_grad():
             for i, batch in enumerate(dataloader):
                 if i == 0:
-                    self.c = self.model.calculate_mean(batch.image.to(self.dev))
+                    self.c = self.model.calculate_mean(batch.image.to(self.device))
                 else:
-                    self.c += self.model.calculate_mean(batch.image.to(self.dev))
+                    self.c += self.model.calculate_mean(batch.image.to(self.device))
 
             self.c /= len(dataloader)
 

src/anomalib/models/image/glass/torch_model.py

@@ -348,8 +348,6 @@ def __init__(
 
         self.augmentor = PerlinAnomalyGenerator(anomaly_source_path)
 
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
         self.focal_loss = FocalLoss()
 
         self.forward_modules = torch.nn.ModuleDict({})
@@ -550,6 +548,7 @@ def forward(
         self,
         img: torch.Tensor,
         c: torch.Tensor | None = None,
+        device: torch.device | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Forward pass to compute patch-wise feature embeddings for original and augmented images.
 
@@ -573,7 +572,7 @@ def forward(
         )
         mask_s_gt = mask_s_resized.reshape(-1, 1)
 
-        noise = torch.normal(0, self.noise, true_feats.shape).to(self.device)
+        noise = torch.normal(0, self.noise, true_feats.shape).to(device)
         gaus_feats = true_feats + noise
 
         center = c.repeat(img.shape[0], 1, 1)
@@ -584,7 +583,7 @@ def forward(
         )
         c_t_points = torch.concat([center[mask_s_gt[:, 0] == 0], center], dim=0)
         dist_t = torch.norm(true_points - c_t_points, dim=1)
-        r_t = torch.tensor([torch.quantile(dist_t, q=self.radius)]).to(self.device)
+        r_t = torch.tensor([torch.quantile(dist_t, q=self.radius)]).to(device)
 
         for step in range(self.step + 1):
             scores = self.discriminator(torch.cat([true_feats, gaus_feats]))