1import torch
 2from fractal_block import TwoChannelFractalBlock
 3
 4device = 'cuda' if torch.cuda.is_available() else 'cpu'
 5try:
 6    m = TwoChannelFractalBlock(8).to(device)
 7    a = torch.randn(4, 3, 8, device=device)
 8    b = torch.randn(4, 3, 8, device=device)
 9    ap, bp = m(a, b)
10    assert ap.shape == bp.shape == (4, 8)
11    assert torch.isfinite(ap).all() and torch.isfinite(bp).all()
12    print({'device': device, 'a_parent': tuple(ap.shape), 'b_parent': tuple(bp.shape), 'coefficients': m.coefficients()})
13except Exception as e:
14    if device == 'cuda':
15        m = TwoChannelFractalBlock(8)
16        ap, bp = m(torch.randn(4,3,8), torch.randn(4,3,8))
17        print({'device': 'cpu-fallback', 'a_parent': tuple(ap.shape), 'b_parent': tuple(bp.shape), 'error': str(e)})
18    else:
19        raise