# Codeblock 6
sam = SAM()
x = torch.randn(1, 512, 28, 28)

out = sam(x)