Fix greyscale injector

2020-09-02 10:29:40 -06:00 · 2020-09-02 10:29:40 -06:00 · d90c96e55e
commit d90c96e55e
parent 8b52d46847
1 changed files with 1 additions and 1 deletions
--- a/codes/models/steps/injectors.py
+++ b/codes/models/steps/injectors.py
@ -101,5 +101,5 @@ class GreyInjector(Injector):

    def forward(self, state):
        mean = torch.mean(state[self.opt['in']], dim=1, keepdim=True)
-        mean = torch.repeat(mean, (-1, 3, -1, -1))
+        mean = mean.repeat((1, 3, 1, 1))
        return {self.opt['out']: mean}