some checks

2022-06-09 21:46:32 -06:00 · 2022-06-09 21:46:32 -06:00 · 07bdd865dc
commit 07bdd865dc
parent 34005367fd
1 changed files with 1 additions and 0 deletions
--- 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