From 088093db5ffc9b5a1d78628b7faa457032a168cb Mon Sep 17 00:00:00 2001
From: chungongyu <chungongyu@gmail.com>
Date: Mon, 30 Mar 2026 14:43:02 +0800
Subject: [PATCH] fix: check weights

---
 profold2/data/dataset.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/profold2/data/dataset.py b/profold2/data/dataset.py
index 4831c02d..2ce16976 100644
--- a/profold2/data/dataset.py
+++ b/profold2/data/dataset.py
@@ -2021,6 +2021,7 @@ def _split_args(args):
   if 'weights' in kwargs:
     weights = kwargs.pop('weights')
     if weights:
+      assert len(dataset) == len(weights)
       kwargs['sampler'] = WeightedRandomSampler(weights, num_samples=len(weights))
       if 'shuffle' in kwargs:
         kwargs.pop('shuffle')