DIOR-VIT: Vision Transformers Learn the Order of Cancer Grades

The problem with treating cancer grades like unrelated boxes

Most deep learning systems for cancer grading are trained the same way an image classifier is trained to tell cats from dogs. Each tissue patch gets a single label, benign, well differentiated, moderately differentiated or poorly differentiated, and the network is scored on whether it picks the right box. That framing works fine for cats and dogs because a cat is not secretly forty percent of the way to being a dog. Cancer grades are different. They sit on a scale. A poorly differentiated tumor is not just a different category from a benign sample, it is further along a spectrum of aggressiveness, and a pathologist’s clinical decision often depends on exactly how far along that spectrum a sample falls.

Jin Tae Kwak’s group at Korea University, working with pathologists Boram Song and Kyungeun Kim at Kangbuk Samsung Hospital and Seegene Medical Foundation, has been chipping away at this mismatch for a few years. An earlier paper from the same broader research line, Le Vuong and colleagues in 2021, showed that adding an ordinal classification task alongside the usual categorical one improves grading accuracy. That paper established the idea that grade order matters. The new paper, published in Medical Image Analysis in 2025 and led by Ju Cheon Lee and Keunho Byeon, pushes the idea further by asking a sharper question. If order matters, does the size of the gap between two grades also carry information the model should learn directly, rather than only knowing that one grade is higher than another.

Order learning already existed, just not for tumors

The idea of comparing two things instead of scoring one thing in isolation is not new. It comes from a line of work called order learning, first applied to guessing a person’s age from a photograph by comparing that photo against a reference face and deciding whether the person looks younger, about the same age, or older. That approach works well for ranking, but the authors point out a real limitation. Order learning only tells you the direction of a comparison, greater than, roughly equal to, or less than. It throws away magnitude. A ten year age gap and a fifty year age gap would both just register as older than. For cancer grading, treating a jump from benign to poorly differentiated the same as a jump from moderately to poorly differentiated would be a mistake, because clinically those two situations are not remotely equivalent.

How DIOR-ViT is built

DIOR-ViT has three working parts bolted onto a shared feature extractor. The feature extractor starts with a ResNet50V2 backbone that turns a tissue patch into a grid of feature maps, which then get sliced into token embeddings and fed through a twelve block transformer encoder, essentially a standard vision transformer with a convolutional front end rather than a raw patch embedding. This produces one feature vector per image, taken from the class token after the final layer normalization.

From that shared feature vector, two heads branch off. A categorical classifier, a small stack of linear and LeakyReLU layers ending in a softmax, predicts which of the grade categories the sample belongs to. That part is conventional. The interesting head is the differential ordinal classifier. During training, every image in a batch is paired with every other image in that batch. For a given pair, the network subtracts one image’s feature vector from the other’s, producing what the authors call a differential feature vector, and feeds that single vector into a one neuron linear layer that predicts how far apart the two samples’ ground truth grades actually are, including the sign.

Formally, for a minuend sample with grade y subscript i and a subtrahend sample with grade y subscript j, the ground truth differential label is defined simply as their difference.

If grades run from zero for benign up to three for the most aggressive category, then comparing a poorly differentiated sample against a benign one gives a target of positive three, while comparing two moderately differentiated samples against each other gives a target near zero. The sign recovers the greater than, less than, or roughly equal comparisons from classic order learning, and the magnitude adds the piece that was previously discarded.

A loss function built to fit the shape of the problem

Once you decide the model should predict a signed real number difference, you need a loss function for that prediction, and this is where the paper spends real effort. Mean squared error and mean absolute error are the default choices for that kind of regression, but the authors note, echoing what Le Vuong’s group found earlier, that plain regression losses do not pair well with a categorical classification task running alongside them. An ordinal cross entropy loss has been proposed before as a fix, but it has two annoyances. It requires converting a distance into a probability through a softmax step before computing cross entropy, and the resulting loss curve is not smooth across the full range of possible distances.

The paper’s answer is a new loss called negative absolute difference log likelihood, NAD for short, which penalizes an incorrect differential prediction on a logarithmic scale without the extra softmax conversion step.

Here P is the set of all sample pairs in a batch, K is the maximum possible gap between grades, and epsilon is a tiny constant to keep the logarithm defined. The categorical loss is standard cross entropy, and the two are combined with a weighting term lambda set to six point five after the authors ran a small search to balance the magnitude of the two losses during early training.

The shape of NAD sits between mean absolute error and mean squared error. Mean absolute error produces a constant gradient regardless of how wrong the prediction is, which is not always ideal because it treats a small mistake and a huge mistake with the same urgency of correction. Mean squared error’s gradient shrinks as the loss approaches zero, which is good for fine convergence but can undershoot when errors are large. NAD keeps a comparatively larger gradient across the range while still smoothing out near a perfect prediction, which the authors argue gives the optimizer a steadier signal throughout training.

What the experiments actually tested

The paper is unusually generous with its evaluation, running three separate tissue types, each with an independent, differently sourced test set to check generalization.

• Colorectal tissue. A public dataset originally assembled by Le Vuong and colleagues, with a training set, a validation set, and two held out test sets. The first test set came from the same tissue microarrays and scanner as training. The second, much larger test set of over one hundred ten thousand images, came from forty five whole slide images collected years later on a different scanner, which is a genuine stress test of generalization.
• Prostate tissue. Two public sources, a Harvard Dataverse dataset scanned with a NanoZoomer scanner, and the external Gleason 2019 challenge dataset scanned with a different Aperio scanner at a different site.
• Gastric tissue. A new dataset the authors themselves contributed, collected from ninety eight whole slide images at Kangbuk Samsung Hospital between 2016 and 2020 under institutional review board approval, reviewed by pathologists Boram Song and Kyungeun Kim, and split into over one hundred fourteen thousand tissue patches across benign and three tumor differentiation grades.

DIOR-ViT was compared against eleven other models spanning three plain convolutional networks, ResNet50, DenseNet121 and EfficientNet B0, a multiscale convolutional model, two prior multitask convolutional models that also used ordinal losses, a self ensemble convolutional model, and four transformer based models including plain ViT, Swin Transformer, DeiT III and a hybrid model called StoHisNet built specifically for gastric pathology.

The headline numbers

On the colorectal dataset’s easier, matched test set, DIOR-ViT reached 87.78 percent accuracy and a quadratic weighted kappa of 0.942, edging out the next best model by roughly half a point of accuracy. The real separation showed up on the harder, out of distribution test set collected years later on different equipment, where DIOR-ViT hit 82.80 percent accuracy against a next best competitor around 77.5 percent, a gap of more than five accuracy points.

A pattern worth sitting with is that DIOR-ViT’s advantage grew on the datasets collected from a different scanner, a different time period, or a different institution than the training data. That is exactly the situation a real deployment faces, since a hospital adopting a grading tool almost never uses the identical scanner and staining protocol the model was trained on. A model that only wins on matched test data is a much shakier proposition than one that holds up under a scanner and site shift, and this paper’s strongest evidence sits precisely there.

Taking the differential task away tells you how much it matters

The ablation table is where the paper makes its strongest case. Strip the differential ordinal classifier out entirely and DIOR-ViT collapses back into an ordinary vision transformer trained with cross entropy alone. That plain version trails the full DIOR-ViT by a consistent margin across every dataset and metric, confirming the differential task is doing real work rather than just adding parameters for their own sake.

Swapping the NAD loss for mean squared error or mean absolute error, while keeping the rest of the architecture identical, also hurts performance, sometimes substantially. On the colorectal datasets, using mean absolute error instead of NAD dropped accuracy by up to about 4.2 percentage points and quadratic weighted kappa by up to 0.043, differences the authors report as statistically meaningful. Swapping in the earlier ordinal cross entropy loss instead of NAD closed some of that gap but still trailed the full NAD version in most settings. This is a reasonably convincing three way comparison, since it isolates the loss function as the variable rather than conflating it with architecture changes.

Where the attention actually looks

The authors ran Grad CAM visualizations comparing plain ViT against DIOR-ViT on the same colorectal images. The pattern they describe is intuitive once you read it. For well differentiated tumor samples, where cancer cells still form recognizable gland structures, DIOR-ViT concentrates attention tightly on the tumor cells forming those glands, while plain ViT scatters its attention across cytoplasm, stroma and open lumen areas somewhat indiscriminately. For moderately differentiated samples, DIOR-ViT shifts focus toward the denser, fibrotic stroma regions associated with tumor invasion and inflammatory infiltration, a detail that lines up with what pathologists actually look for when grading. For poorly differentiated samples, where architecture has broken down almost entirely, DIOR-ViT spreads attention across the whole sample, which the authors attribute to the fact that atypical, irregular cells are distributed everywhere rather than confined to a structure. None of this proves the model reasons the way a pathologist reasons, but it is a reasonable qualitative signal that the differential training objective nudges attention toward histologically relevant regions rather than incidental texture.

Does the differential head actually learn something coherent

Beyond raw accuracy, the authors checked whether the predicted differential values behave the way you would want a genuine ordinal measurement to behave. They averaged predictions across many minuend and subtrahend pairs for every combination of grade categories. Two properties held up. First, larger true gaps between grades produced larger predicted gaps, so the network is not just guessing signs. Second, and more interesting, the predictions were roughly additive. On the prostate datasets, the predicted gap between Gleason grade 5 and grade 3 was close to the sum of the predicted gap between grade 5 and grade 4 plus the predicted gap between grade 4 and grade 3. That additive consistency is a nontrivial property for a network trained only on pairwise comparisons to discover on its own, and it is the strongest evidence in the paper that the differential ordinal classifier learned something resembling a real numeric scale rather than a shortcut.

What kind of backbone the differential idea needs

A separate set of experiments swapped the ResNet plus transformer backbone for other feature extractors, including Swin Transformer and two backbones pretrained specifically on pathology images at large scale, HIPT and Phikon. The ImageNet pretrained ViT backbone used in the main DIOR-ViT model actually outperformed both pathology specific backbones in most settings, which is a mildly surprising result the authors are upfront about. Their explanation is architectural rather than data related. HIPT was built around a hierarchical structure meant for whole slide analysis across cellular, patch and region scales, and using only its lowest level patch representation likely does not showcase its real strength. Phikon, built for patch level representation without that hierarchical constraint, paired better with the differential ordinal objective. The authors are careful to flag that they fully fine tuned every backbone rather than trying more parameter efficient tuning strategies, and that a different tuning recipe might change these particular rankings.

Clinical translation gap

It is worth being direct about the distance between a strong benchmark result and a tool a hospital could actually use. This paper evaluates DIOR-ViT on fixed image patches carved out of whole slide images, not on the full diagnostic workflow of scanning, cropping regions of interest, aggregating patch level predictions into a single case level grade, and presenting that grade with enough context for a pathologist to sign off on it. The gastric dataset, while a genuinely valuable contribution at over one hundred thousand images, comes from a single hospital and ninety eight patients, which is a modest patient count for a diagnostic tool that would eventually need to generalize across demographics, staining protocols and scanner vendors well beyond what one institution can represent. The colorectal and prostate results do include cross institution test sets, which is a real strength, but even those are limited to two source institutions each rather than the dozens a national deployment would eventually touch. None of this diminishes the technical contribution, but it does mean the honest next step is prospective, multi site validation before anything resembling clinical decision support, not a leap from patch level accuracy straight to bedside use.

Clinical limitations reported in the paper

• The gastric dataset was collected from a single hospital and ninety eight patients, which limits how confidently its results generalize to other populations or scanning equipment.
• All datasets classify pre extracted tissue patches rather than performing end to end whole slide image analysis, so the paper does not test the full pipeline a deployed tool would need, including region of interest selection and case level aggregation.
• The prostate and colorectal cross institution test sets came from only two sources each, so the demonstrated robustness, while real, covers a narrow slice of the scanner and protocol diversity a wider rollout would encounter.
• The paper reports patch level classification metrics only. It does not report how grading errors at the patch level would translate into case level diagnostic errors that affect patient management.

Where this sits relative to the field

Multitask learning that pairs a categorical loss with some form of ordinal signal is not a new idea in computational pathology, and the authors are careful to credit the 2021 Le Vuong paper this work builds on directly. What is new here is combining that ordinal signal with the pairwise comparison structure of order learning, which previously had only been applied to problems like age estimation from faces. The authors position DIOR-ViT as, to their knowledge, the first attempt to merge these two lines of research inside computational pathology, and the additive consistency result described earlier is reasonable support for that claim holding real substance rather than being a cosmetic combination.

It is also worth noting the model complexity story is more favorable than the extra pairwise comparison machinery might suggest. Because the feature extractor still only processes one image at a time and the pairwise subtraction happens after that, DIOR-ViT trains and runs inference at nearly the same speed as a plain vision transformer, despite technically doing far more comparisons during training. That efficiency argument matters for anyone considering this approach for a real dataset, since order learning style pairwise methods have a reputation for being expensive to train.

Limitations worth taking seriously

Beyond the clinical translation concerns above, a few technical caveats stand out from a close read. The weighting term lambda that balances the two loss functions was tuned empirically on one dataset, the colorectal training set, and then reused across prostate and gastric experiments. The ablation on lambda shows performance is somewhat sensitive to this choice, and the paper itself notes the optimal value likely depends on the dataset. A team applying this method to a new tissue type would probably need to redo that tuning rather than trust the reported six point five value out of the box. The ordinal cross entropy loss used as a comparison baseline is attributed to a 2020 object detection paper by Carion and colleagues in the references, which reads as a possible citation mismatch worth the authors clarifying, since ordinal cross entropy losses for grading tasks are more commonly associated with the earlier Le Vuong line of work this paper otherwise builds on directly. Readers relying on that specific baseline attribution should treat it cautiously until clarified. Finally, all three tissue types use grading schemes with only four ordinal classes. It remains an open question whether the differential approach continues to help as cleanly on cancers graded on finer scales, such as full Gleason scores rather than grade groups, where the number of possible pairwise gaps grows substantially.

Conclusion

The core achievement of this paper is a fairly narrow but genuinely useful idea, executed carefully across three different cancer types with real cross institution testing rather than a single convenient benchmark. Cancer grades are not just labels, they carry a magnitude of difference that matters clinically, and DIOR-ViT is a concrete demonstration that a network trained to predict that magnitude directly, through pairwise comparisons and a custom loss function, ends up with a more useful feature space than one trained on category labels alone.

The conceptual shift underneath the engineering is worth sitting with on its own. Most ordinal classification work in medical imaging still treats order as a constraint to respect, something to prevent the model from predicting wildly out of sequence grades. This paper treats order as a training signal in its own right, something the model can actively learn from by comparing samples against each other rather than only against a fixed label. That reframing, from order as a rule to order as data, is the part likely to transfer well beyond cancer grading.

Transferability seems genuinely plausible here. Any diagnostic or severity scale with a natural ordering, disease staging, retinal disease severity, fibrosis scoring, could in principle benefit from the same differential comparison trick, and the authors say as much in their own framing of the contribution. The additive consistency result on the prostate dataset is the most persuasive piece of evidence that the learned differences behave like a real numeric scale rather than a superficial pattern match, which is exactly the property you would want before trusting this kind of signal in another domain.

The honest remaining limitations are the ones any careful reader should hold onto. A single hospital gastric dataset, patch level rather than whole slide evaluation, a hand tuned loss weighting term, and a citation detail worth double checking all temper how far the current results should be extrapolated. None of these undercut the central finding, but they do mean the path from this paper to a bedside tool runs through considerably more validation, not around it.

Where this heads next seems fairly clear from the authors’ own stated plans. They flag whole slide image level classification, where patches get aggregated into a single case level prediction through multiple instance learning, as the natural extension, along with exploring parameter efficient fine tuning for the pathology specific backbones that underperformed in this study. If differential ordinal learning holds up as cleanly at the whole slide level as it does at the patch level here, it would be a meaningfully useful addition to the computational pathology toolkit, not because it reinvents the vision transformer, but because it finally gives the model a way to know not just what grade a tumor is, but how much that grade actually differs from the alternative.

Frequently asked questions

What does differential ordinal learning mean in plain terms

It means training a model to predict not just a single sample’s category but the gap between two samples’ categories, including how large that gap is and which direction it runs. For cancer grading that lets the model learn that the jump from benign to aggressive tumor is much larger than the jump between two adjacent tumor grades.

Is DIOR-ViT the same as a standard vision transformer

The feature extraction backbone is a fairly standard vision transformer built on top of a ResNet50V2 front end. What makes DIOR-ViT different is the second prediction head, the differential ordinal classifier, and the new NAD loss function used to train it alongside the usual categorical classifier.

Did this improve accuracy on every single dataset tested

DIOR-ViT achieved the best accuracy and quadratic weighted kappa across nearly every test set in the paper, though on two of the five test sets a different multitask model edged it out narrowly on the F1 macro metric specifically, so the win is broad but not universally the single best score on every metric.

Has this model been tested in an actual hospital workflow

No. The evaluation is on pre extracted tissue image patches from research datasets, not on a live clinical pipeline, and the paper does not report prospective validation in a hospital setting. That step remains future work.

Why does the NAD loss function matter instead of just using mean squared error

The ablation experiments show that swapping NAD for mean squared error or mean absolute error measurably lowers accuracy and quadratic weighted kappa across every dataset tested, so the specific shape of the loss curve, not just the presence of an ordinal task, contributes to the final performance.

Could this approach work for diseases other than cancer

The authors argue it should generalize to any diagnostic problem with ordered severity labels, such as disease staging or symptom severity scoring, though the paper itself only tests cancer grading in colorectal, prostate and gastric tissue, so that broader claim has not yet been demonstrated experimentally.

Reproducible PyTorch implementation

The block below is an independent, from scratch implementation of the architecture and loss function described in the paper, written to match the equations and design choices above. It is meant as a starting point for experimentation, not a copy of the authors’ original codebase, which was not released publicly at the time of writing.

# dior_vit.py
# Independent reproduction of DIOR-ViT (Lee et al., Medical Image Analysis 2025)
# Categorical classification + differential ordinal classification with NAD loss

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import math


class PositionalEmbedding2D(nn.Module):
    """Sin cos positional embedding added to the token sequence, plus a class token."""

    def __init__(self, num_tokens, dim):
        super().__init__()
        pe = torch.zeros(num_tokens, dim)
        position = torch.arange(0, num_tokens).unsqueeze(1).float()
        div_term = torch.exp(torch.arange(0, dim, 2).float() * (-math.log(10000.0) / dim))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x):
        return x + self.pe[:, : x.size(1)]


class FeatureExtractor(nn.Module):
    """ResNet50V2 style stem, linear token projection, class token, positional embedding,
    then a standard twelve block transformer encoder."""

    def __init__(self, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0):
        super().__init__()
        resnet = models.resnet50(weights=None)
        self.stem = nn.Sequential(*list(resnet.children())[:-2])  # -> (B, 2048, 24, 24) for 384 input
        self.proj = nn.Conv2d(2048, embed_dim, kernel_size=1)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = PositionalEmbedding2D(577, embed_dim)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            nhead=num_heads,
            dim_feedforward=int(embed_dim * mlp_ratio),
            dropout=0.1,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=depth)
        self.norm = nn.LayerNorm(embed_dim)

        nn.init.trunc_normal_(self.cls_token, std=0.02)

    def forward(self, x):
        feats = self.stem(x)                      # (B, 2048, H, W)
        feats = self.proj(feats)                   # (B, D, H, W)
        B, D, H, W = feats.shape
        tokens = feats.flatten(2).transpose(1, 2)  # (B, H*W, D)
        cls = self.cls_token.expand(B, -1, -1)
        tokens = torch.cat([cls, tokens], dim=1)
        tokens = self.pos_embed(tokens)
        encoded = self.encoder(tokens)
        cls_out = self.norm(encoded[:, 0])          # (B, D) feature vector f
        return cls_out


class CategoricalClassifier(nn.Module):
    def __init__(self, embed_dim=768, num_classes=4, hidden=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(embed_dim, hidden),
            nn.LeakyReLU(0.1),
            nn.Linear(hidden, hidden),
            nn.LeakyReLU(0.1),
            nn.Linear(hidden, num_classes),
        )

    def forward(self, f):
        return self.net(f)  # logits, softmax applied in the loss


class DifferentialOrdinalClassifier(nn.Module):
    """Single linear layer that maps a differential feature vector f_i - f_j to a scalar."""

    def __init__(self, embed_dim=768):
        super().__init__()
        self.linear = nn.Linear(embed_dim, 1)

    def forward(self, f_diff):
        return self.linear(f_diff).squeeze(-1)


class DIORViT(nn.Module):
    def __init__(self, num_classes=4, embed_dim=768):
        super().__init__()
        self.feature_extractor = FeatureExtractor(embed_dim=embed_dim)
        self.categorical_head = CategoricalClassifier(embed_dim, num_classes)
        self.differential_head = DifferentialOrdinalClassifier(embed_dim)

    def forward(self, x_batch):
        """x_batch: (B, 3, H, W). Returns categorical logits and, for every ordered
        pair in the batch, the predicted differential label."""
        f = self.feature_extractor(x_batch)          # (B, D)
        logits = self.categorical_head(f)             # (B, C)

        B = f.size(0)
        f_i = f.unsqueeze(1).expand(B, B, -1)          # (B, B, D)
        f_j = f.unsqueeze(0).expand(B, B, -1)          # (B, B, D)
        f_diff = (f_i - f_j).reshape(B * B, -1)        # (B*B, D)
        r_hat = self.differential_head(f_diff).reshape(B, B)  # (B, B), diagonal unused

        return logits, r_hat, f


def nad_loss(r_hat, r_true, K, eps=1e-5, mask=None):
    """Negative absolute difference log likelihood loss, equation 11 in the paper.

    r_hat, r_true: (B, B) predicted and ground truth differential labels
    K: |c_max - c_min|, the maximum possible gap between class labels
    mask: optional (B, B) boolean tensor, True for valid off diagonal pairs
    """
    if mask is None:
        B = r_hat.size(0)
        mask = ~torch.eye(B, dtype=torch.bool, device=r_hat.device)

    diff = torch.abs(r_true - r_hat)
    inside = 1.0 - diff / (2.0 * K + eps)
    inside = torch.clamp(inside, min=eps)  # guard the log against non positive input
    per_pair_loss = -torch.log(inside)

    per_pair_loss = per_pair_loss[mask]
    return per_pair_loss.mean()


def dior_vit_total_loss(logits, y_true, r_hat, K, lam=6.5):
    """Combines the categorical cross entropy loss with lambda times the NAD loss."""
    B = logits.size(0)
    y_i = y_true.unsqueeze(1).expand(B, B).float()
    y_j = y_true.unsqueeze(0).expand(B, B).float()
    r_true = y_i - y_j  # ground truth differential labels, equation 8

    cat_loss = F.cross_entropy(logits, y_true)
    diff_loss = nad_loss(r_hat, r_true, K)
    total = cat_loss + lam * diff_loss
    return total, cat_loss, diff_loss


def train_one_epoch(model, loader, optimizer, device, num_classes=4, lam=6.5):
    model.train()
    K = float(num_classes - 1)
    running_loss = 0.0

    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()

        logits, r_hat, _ = model(images)
        loss, cat_loss, diff_loss = dior_vit_total_loss(logits, labels, r_hat, K, lam)

        loss.backward()
        optimizer.step()
        running_loss += loss.item() * images.size(0)

    return running_loss / len(loader.dataset)


@torch.no_grad()
def evaluate(model, loader, device, num_classes=4):
    model.eval()
    correct, total = 0, 0

    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        logits, _, _ = model(images)
        preds = logits.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)

    accuracy = correct / max(total, 1)
    return {"accuracy": accuracy}


def smoke_test():
    """Runs one forward pass, one backward pass, and one evaluation step on random
    dummy data, just to confirm every shape lines up before real training."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    num_classes = 4
    batch_size = 6

    model = DIORViT(num_classes=num_classes).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    dummy_images = torch.randn(batch_size, 3, 384, 384, device=device)
    dummy_labels = torch.randint(0, num_classes, (batch_size,), device=device)

    logits, r_hat, f = model(dummy_images)
    assert logits.shape == (batch_size, num_classes)
    assert r_hat.shape == (batch_size, batch_size)
    assert f.shape[0] == batch_size

    K = float(num_classes - 1)
    loss, cat_loss, diff_loss = dior_vit_total_loss(logits, dummy_labels, r_hat, K, lam=6.5)
    loss.backward()
    optimizer.step()

    print(f"smoke test passed, total loss {loss.item():.4f}, "
          f"cat loss {cat_loss.item():.4f}, diff loss {diff_loss.item():.4f}")


if __name__ == "__main__":
    smoke_test()

Related reading