From b49df0a1ed4477707de8b4d2f204865d49f92f5a Mon Sep 17 00:00:00 2001
From: Thien Tran <gau.nernst@yahoo.com.sg>
Date: Tue, 8 Aug 2023 21:51:18 +0800
Subject: [PATCH] add sanity check

---
 vision_toolbox/backbones/vit.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/vision_toolbox/backbones/vit.py b/vision_toolbox/backbones/vit.py
index 829cf21..ae54506 100644
--- a/vision_toolbox/backbones/vit.py
+++ b/vision_toolbox/backbones/vit.py
@@ -81,6 +81,7 @@ def __init__(
         norm: _norm = partial(nn.LayerNorm, eps=1e-6),
         act: _act = nn.GELU,
     ) -> None:
+        assert img_size % patch_size == 0
         super().__init__()
         self.patch_embed = nn.Conv2d(3, d_model, patch_size, patch_size)
         self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model)) if cls_token else None