Upload model

modeling_basnet.py

@@ -14,7 +14,7 @@ logger = logging.getLogger(__name__)
 
 
 @dataclass
-class BASNetModelOutput(ModelOutput):
+class BasNetSideOutput(ModelOutput):
     dout: torch.Tensor
     d1: Optional[torch.Tensor] = None
     d2: Optional[torch.Tensor] = None
@@ -25,6 +25,11 @@ class BASNetModelOutput(ModelOutput):
     db: Optional[torch.Tensor] = None
 
 
+@dataclass
+class BASNetModelOutput(ModelOutput):
+    activated: BasNetSideOutput
+
+
 class RefUnet(nn.Module):
     def __init__(self, in_ch: int, inc_ch: int) -> None:
         super().__init__()
@@ -466,27 +471,21 @@ class BASNetModel(PreTrainedModel):
         d6_act = torch.sigmoid(d6)
         db_act = torch.sigmoid(db)
 
+        side_outputs = (
+            dout_act,
+            d1_act,
+            d2_act,
+            d3_act,
+            d4_act,
+            d5_act,
+            d6_act,
+            db_act,
+        )
         if not return_dict:
-            return (
-                dout_act,
-                d1_act,
-                d2_act,
-                d3_act,
-                d4_act,
-                d5_act,
-                d6_act,
-                db_act,
-            )
+            return (side_outputs,)
 
         return BASNetModelOutput(
-            dout=dout_act,
-            d1=d1_act,
-            d2=d2_act,
-            d3=d3_act,
-            d4=d4_act,
-            d5=d5_act,
-            d6=d6_act,
-            db=db_act,
+            activated=BasNetSideOutput(*side_outputs),
         )