cast np.int32 as int for random.Random arg

mlcommons · Jan 23, 2025 · 5775ed1 · 5775ed1
1 parent 9b7cee4
commit 5775ed1
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/algorithmic_efficiency/workloads/cifar/cifar_pytorch/workload.py b/algorithmic_efficiency/workloads/cifar/cifar_pytorch/workload.py
@@ -82,7 +82,7 @@ def _build_dataset(
     }
     if split == 'eval_train':
       train_indices = indices_split['train']
-      random.Random(data_rng[0]).shuffle(train_indices)
+      random.Random(int(data_rng[0])).shuffle(train_indices)
       indices_split['eval_train'] = train_indices[:self.num_eval_train_examples]
     if split in indices_split:
       dataset = torch.utils.data.Subset(dataset, indices_split[split])

diff --git a/algorithmic_efficiency/workloads/imagenet_resnet/imagenet_pytorch/workload.py b/algorithmic_efficiency/workloads/imagenet_resnet/imagenet_pytorch/workload.py
@@ -120,7 +120,7 @@ def _build_dataset(
 
     if split == 'eval_train':
       indices = list(range(self.num_train_examples))
-      random.Random(data_rng[0]).shuffle(indices)
+      random.Random(int(data_rng[0])).shuffle(indices)
       dataset = torch.utils.data.Subset(dataset,
                                         indices[:self.num_eval_train_examples])
 

diff --git a/algorithmic_efficiency/workloads/librispeech_conformer/librispeech_pytorch/workload.py b/algorithmic_efficiency/workloads/librispeech_conformer/librispeech_pytorch/workload.py
@@ -166,7 +166,7 @@ def _build_input_queue(
     ds = LibriSpeechDataset(split=ds_split, data_dir=data_dir)
     if split == 'eval_train':
       indices = list(range(len(ds)))
-      random.Random(data_rng[0]).shuffle(indices)
+      random.Random(int(data_rng[0])).shuffle(indices)
       ds = torch.utils.data.Subset(ds, indices[:self.num_eval_train_examples])
 
     sampler = None