12 lines
342 B
Python
12 lines
342 B
Python
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')
|