Thanks for the code!
It looks generally good, but this initialization should be problematic:
self.heads = [AttentionHead(2, 1) for _ in range(config.num_heads)]

Creating layers in a plain Python list will not properly register their parameters in the parent module and thus neither will model.parameters() return them (i.e. the optimizer might not update them) nor will model.to('cuda') work on them (and they will stay on the same device).
Use nn.ModuleList instead and rerun the code.
Let me know, if this works.