fix unit test

openvinotoolkit · Dec 18, 2024 · 045e3ed · 045e3ed
1 parent b44ebea
commit 045e3ed
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/otx/algo/common/losses/cross_focal_loss.py b/src/otx/algo/common/losses/cross_focal_loss.py
@@ -7,7 +7,6 @@
 
 import torch
 import torch.nn.functional
-from otx.utils.device import get_available_device
 from torch import Tensor, nn
 from torch.cuda.amp import custom_fwd
 
@@ -80,7 +79,7 @@ def __init__(
 
         self.cls_criterion = cross_sigmoid_focal_loss
 
-    @custom_fwd(device_type=get_available_device(), cast_inputs=torch.float32)
+    @custom_fwd(cast_inputs=torch.float32)
     def forward(
         self,
         pred: Tensor,