add args

Signed-off-by: tangy5 <[email protected]>
Project-MONAI · Apr 5, 2023 · cb1402c · cb1402c
1 parent 9923686
commit cb1402c
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/monai/networks/nets/universal_model.py b/monai/networks/nets/universal_model.py
@@ -102,6 +102,16 @@ def __init__(
         img_size: Union[Sequence[int], int],
         in_channels: int,
         out_channels: int,
+        depths: Sequence[int] = (2, 2, 2, 2),
+        num_heads: Sequence[int] = (3, 6, 12, 24),
+        feature_size: int = 48,
+        norm_name: Union[Tuple, str] = "instance",
+        drop_rate: float = 0.0,
+        attn_drop_rate: float = 0.0,
+        dropout_path_rate: float = 0.0,
+        normalize: bool = True,
+        use_checkpoint: bool = False,
+        spatial_dims: int = 3,
     ):
         super(SwinUNETR_backbone, self).__init__(img_size,in_channels,out_channels,feature_size=48)