1 files changed, 4 insertions, 4 deletions

diff --git a/util/noise.py b/util/noise.py
index 3c4f82d..e3ebdb2 100644
--- a/util/noise.py
+++ b/util/noise.py
@@ -48,13 +48,13 @@ def rand_perlin_2d_octaves(shape, res, octaves=1, persistence=0.5, dtype=None, d
     return noise
 
 
-def perlin_noise(batch_size: int, channels: int, width: int, height: int, res=8, octaves=1, dtype=None, device=None, generator=None):
+def perlin_noise(shape: tuple[int, int, int, int], res=8, octaves=1, dtype=None, device=None, generator=None):
     return torch.stack([
         torch.stack([
             rand_perlin_2d_octaves(
-                (width, height), (res, res), octaves, dtype=dtype, device=device, generator=generator
+                (shape[2], shape[3]), (res, res), octaves, dtype=dtype, device=device, generator=generator
             )
-            for _ in range(channels)
+            for _ in range(shape[1])
         ])
-        for _ in range(batch_size)
+        for _ in range(shape[0])
     ])