File size: 312 Bytes
45eaa9d | 1 2 3 4 5 6 7 8 9 | import torch
from hilbert_mi.scalable import HilbertOperatorClassifier
def test_forward_and_backward():
m=HilbertOperatorClassifier(width=8,depth=1,modes=4,rank=2)
x=torch.randn(2,3,16,16); y=m(x)
assert y.shape==(2,10); y.sum().backward()
assert all(p.grad is not None for p in m.parameters())
|