While users are expected to follow the constructor definition of VisionTransformer, the current handling of norm_layer can be misleading.

In flash_attn/models/vit.py:164, the code sets:

norm_layer = norm_layer or functools.partial(nn.LayerNorm, eps=1e-6)

This means that if a user passes norm_layer=False, Python treats it as falsy and the model still uses the default LayerNorm.

Minimal Repro

import torch
import torch.nn as nn
from flash_attn.models.vit import VisionTransformer

m = VisionTransformer(
    img_size=32, patch_size=16, in_chans=3,
    num_classes=10, embed_dim=64, depth=2, num_heads=4,
    norm_layer=False,   # user may expect to disable normalization
)

print("norm1 type:", type(m.blocks[0].norm1))

Output:

norm1 type: <class 'torch.nn.modules.normalization.LayerNorm'>

Why this is confusing

Even though the code implies norm_layer should be a callable (e.g. nn.LayerNorm), users may assume that passing False will disable normalization, as is common in other APIs (bias=False, norm=False, etc.). Instead, their input is silently ignored and LayerNorm is applied.

Suggested improvement

To reduce confusion, the constructor could validate norm_layer explicitly:

if norm_layer is None:
    norm_layer = functools.partial(nn.LayerNorm, eps=1e-6)
elif not callable(norm_layer):
    raise TypeError(
        f"norm_layer must be a callable (e.g., nn.LayerNorm) or None, got {type(norm_layer)}"
    )

This would prevent silent overrides and give clear feedback if the user passes an unsupported value like False.