diff --git a/vit_pytorch/simple_vit_with_register_tokens.py b/vit_pytorch/simple_vit_with_register_tokens.py
index 6b5ab44..774ee85 100644
--- a/vit_pytorch/simple_vit_with_register_tokens.py
+++ b/vit_pytorch/simple_vit_with_register_tokens.py
@@ -1,3 +1,8 @@
+"""
+    Vision Transformers Need Registers
+    https://arxiv.org/abs/2309.16588
+"""
+
 import torch
 from torch import nn