Update depth_anything_v2/dinov2_layers/attention.py

depth_anything_v2/dinov2_layers/attention.py

@@ -64,18 +64,16 @@ class Attention(nn.Module):
 
 class MemEffAttention(Attention):
     def forward(self, x: Tensor, attn_bias=None) -> Tensor:
-        if not XFORMERS_AVAILABLE:
+        # If xformers is not available, or input is not CUDA, use vanilla attention
+        if (not XFORMERS_AVAILABLE) or (not x.is_cuda):
             assert attn_bias is None, "xFormers is required for nested tensors usage"
             return super().forward(x)
-
+        # Otherwise, use memory efficient attention
         B, N, C = x.shape
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
-
         q, k, v = unbind(qkv, 2)
-
         x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
         x = x.reshape([B, N, C])
-
         x = self.proj(x)
         x = self.proj_drop(x)
         return x