import torch
from DCNv4 import DCNv4

dcn = DCNv4(channels=64, group=4).cuda().eval()
print('--- 50 forward calls, no_grad ---')
with torch.no_grad():
    for s in range(50):
        x = torch.randn(8, 4096, 64, device='cuda')
        _ = dcn(x)
        if s % 10 == 0:
            print(f'{s}: {torch.cuda.memory_allocated() / 1e6:.1f} MB')
