lc

2025-06-04 18:24:17 +08:00
parent 56c9e91a3b
commit a7a28891b1
1 changed files with 4 additions and 1 deletions
@@ -315,7 +315,10 @@ class SmallUnetV5(nn.Module):


 def get_posembed_linear(length, dim, dtype, device):
-    return get_1d_sincos_pos_embed_from_grid(dim, torch.arange(length), output_type='pt').to(dtype=dtype, device=device)
+    try:
+        return get_1d_sincos_pos_embed_from_grid(dim, torch.arange(length), output_type='pt').to(dtype=dtype, device=device)
+    except:
+        return get_1d_sincos_pos_embed_from_grid(dim, torch.arange(length)).to(dtype=dtype, device=device)


 class STKAttentionV5(nn.Module):