Tan_pytorch_segmentation/pytorch_segmentation/PV_Attention/CoTAttention.py

13 lines
233 B
Python

from attention.CoTAttention import CoTAttention
import torch
from torch import nn
from torch.nn import functional as F
input=torch.randn(50,512,7,7)
cot = CoTAttention(dim=512,kernel_size=3)
output=cot(input)
print(output.shape)