diff --git a/src/models/srnet.py b/src/models/srnet.py
index 373e160..df74b5e 100644
--- a/src/models/srnet.py
+++ b/src/models/srnet.py
@@ -443,16 +443,16 @@ class SRMsbLsbFlipNet(SRNetBase):
         self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
         self.flip_functions = [
             lambda x: x,
-            lambda x: x[:,:,::-1,:],
-            lambda x: x[:,:,:,::-1],
-            lambda x: x[:,:,::-1,::-1],
+            lambda x: torch.flip(x, dims=[-2]),
+            lambda x: torch.flip(x, dims=[-1]),
+            lambda x: torch.flip(x, dims=[-2, -1]),
         ]
             
     def forward(self, x, config=None):
         b,c,h,w = x.shape
         x = x.reshape(b*c, 1, h, w)
         output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
-        for flip_f in self.flips_functions:
+        for flip_f in self.flip_functions:
            fliped_x = flip_f(x)
            fliped_lsb = fliped_x % 16
            fliped_msb = fliped_x - fliped_lsb           
@@ -461,7 +461,7 @@ class SRMsbLsbFlipNet(SRNetBase):
            if not config is None and config.current_iter % config.display_step == 0:
               config.writer.add_histogram('output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)
               config.writer.add_histogram('output_msb', output_msb.detach().cpu().numpy(), config.current_iter)
-           output += torch.rot90(output_msb + output_lsb, k=-rotations_count, dims=[2, 3])
+           output += flip_f(output_msb + output_lsb)
         output /= 4
         x = output
         x = x.reshape(b, c, h*self.scale, w*self.scale)