CPU edits

birefnet.py

@@ -383,7 +383,7 @@ class PyramidVisionTransformerImpr(nn.Module):
                                               embed_dim=embed_dims[3])
 
         # transformer encoder
-        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]        
         cur = 0
         self.block1 = nn.ModuleList([Block(
             dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale,
@@ -1128,8 +1128,12 @@ class SwinTransformer(nn.Module):
         self.pos_drop = nn.Dropout(p=drop_rate)
 
         # stochastic depth
-        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
-
+        # stochastic depth decay rule (pure python: safe even if model is being initialized on `meta`)
+        total_depth = int(sum(depths))
+        if total_depth <= 1:
+            dpr = [0.0]
+        else:
+            dpr = [float(drop_path_rate) * i / (total_depth - 1) for i in range(total_depth)]
         # build layers
         self.layers = nn.ModuleList()
         for i_layer in range(self.num_layers):