1 files changed, 9 insertions, 6 deletions

diff --git a/util/noise.py b/util/noise.py
index 38ab172..3c4f82d 100644
--- a/util/noise.py
+++ b/util/noise.py
@@ -1,4 +1,5 @@
 import math
+
 import torch
 
 # 2D Perlin noise in PyTorch https://gist.github.com/vadimkantorov/ac1b097753f217c5c11bc2ff396e0a57
@@ -47,11 +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, width: int, height: int, res=8, octaves=1, dtype=None, device=None, generator=None):
+def perlin_noise(batch_size: int, channels: int, width: int, height: int, res=8, octaves=1, dtype=None, device=None, generator=None):
     return torch.stack([
-        rand_perlin_2d_octaves(
-            (width, height), (res, res), octaves, dtype=dtype, device=device, generator=generator
-        ).unsqueeze(0)
-        for _
-        in range(batch_size)
+        torch.stack([
+            rand_perlin_2d_octaves(
+                (width, height), (res, res), octaves, dtype=dtype, device=device, generator=generator
+            )
+            for _ in range(channels)
+        ])
+        for _ in range(batch_size)
     ])