diff --git a/codes/trainer/injectors/gaussian_diffusion_injector.py b/codes/trainer/injectors/gaussian_diffusion_injector.py
index 3d71e2b9..0e35c361 100644
--- a/codes/trainer/injectors/gaussian_diffusion_injector.py
+++ b/codes/trainer/injectors/gaussian_diffusion_injector.py
@@ -21,6 +21,7 @@ def masked_channel_balancer(inp, proportion=1):
 
 
 def channel_restriction(inp, low, high):
+    assert low > 0 and low < inp.shape[1] and high <= inp.shape[1]
     m = torch.zeros_like(inp)
     m[:,low:high] = 1
     return inp * m