From 6cbde80d2ab16ebd0b0aeeb7544da7c370f44aee Mon Sep 17 00:00:00 2001
From: Gao Enhao <gaoeh@lamda.nju.edu.cn>
Date: Sun, 11 Jun 2023 22:42:03 +0800
Subject: [PATCH] [ENH] add InfiniteSampler

---
 examples/hed/utils.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/examples/hed/utils.py b/examples/hed/utils.py
index cf35eaf..42b7316 100644
--- a/examples/hed/utils.py
+++ b/examples/hed/utils.py
@@ -1,6 +1,21 @@
 import torch
 import torch.nn as nn
 import numpy as np
+import torch.utils.data.sampler as sampler
+
+
+class InfiniteSampler(sampler.Sampler):
+    def __init__(self, num_samples):
+        self.num_samples = num_samples
+
+    def __iter__(self):
+        while True:
+            order = np.random.permutation(self.num_samples)
+            for i in range(self.num_samples):
+                yield order[i]
+
+    def __len__(self):
+        return None
 
 
 def gen_mappings(chars, symbs):