more stochastic depth stuff

Author: gau-nernst

timm/models/maxxvit.py

@@ -1174,7 +1174,7 @@ def __init__(
 
         num_stages = len(cfg.embed_dim)
         assert len(cfg.depths) == num_stages
-        dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(cfg.depths)).split(cfg.depths)]
+        dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(cfg.depths), device="cpu").split(cfg.depths)]
         in_chs = self.stem.out_chs
         stages = []
         for i in range(num_stages):

timm/models/mvitv2.py

@@ -749,7 +749,7 @@ def __init__(
         num_stages = len(cfg.embed_dim)
         feat_size = patch_dims
         curr_stride = max(cfg.patch_stride)
-        dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(cfg.depths)).split(cfg.depths)]
+        dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(cfg.depths), device="cpu").split(cfg.depths)]
         self.stages = nn.ModuleList()
         self.feature_info = []
         for i in range(num_stages):

timm/models/nfnet.py

@@ -337,7 +337,7 @@ def __init__(
         )
 
         self.feature_info = [stem_feat]
-        drop_path_rates = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(cfg.depths)).split(cfg.depths)]
+        drop_path_rates = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(cfg.depths), device="cpu").split(cfg.depths)]
         prev_chs = stem_chs
         net_stride = stem_stride
         dilation = 1

timm/models/pit.py

@@ -183,7 +183,7 @@ def __init__(
 
         transformers = []
         # stochastic depth decay rule
-        dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(depth)).split(depth)]
+        dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(depth), device="cpu").split(depth)]
         prev_dim = embed_dim
         for i in range(len(depth)):
             pool = None

timm/models/rdnet.py

@@ -213,7 +213,7 @@ def __init__(
         self.num_stages = len(growth_rates)
         curr_stride = stem_stride
         num_features = num_init_features
-        dp_rates = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(num_blocks_list)).split(num_blocks_list)]
+        dp_rates = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(num_blocks_list), device="cpu").split(num_blocks_list)]
 
         dense_stages = []
         for i in range(self.num_stages):

timm/models/resnetv2.py

@@ -467,7 +467,7 @@ def __init__(
         prev_chs = stem_chs
         curr_stride = 4
         dilation = 1
-        block_dprs = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(layers)).split(layers)]
+        block_dprs = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(layers), device="cpu").split(layers)]
         if preact:
             block_fn = PreActBasic if basic else PreActBottleneck
         else:

timm/models/tresnet.py

@@ -135,7 +135,7 @@ def __init__(
             self.inplanes = self.inplanes // 8 * 8
             self.planes = self.planes // 8 * 8
 
-        dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(layers)).split(layers)]
+        dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(layers), device="cpu").split(layers)]
         conv1 = ConvNormAct(in_chans * 16, self.planes, stride=1, kernel_size=3, act_layer=act_layer)
         layer1 = self._make_layer(
             Bottleneck if v2 else BasicBlock,

timm/models/vovnet.py

@@ -213,7 +213,7 @@ def __init__(
         current_stride = stem_stride
 
         # OSA stages
-        stage_dpr = torch.split(torch.linspace(0, drop_path_rate, sum(block_per_stage)), block_per_stage)
+        stage_dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(block_per_stage), device="cpu").split(block_per_stage)]
         in_ch_list = stem_chs[-1:] + stage_out_chs[:-1]
         stage_args = dict(residual=cfg["residual"], depthwise=cfg["depthwise"], attn=cfg["attn"], **conv_kwargs)
         stages = []