diff --git a/src/models/srlut.py b/src/models/srlut.py
index ed80dda..11bee80 100644
--- a/src/models/srlut.py
+++ b/src/models/srlut.py
@@ -4,7 +4,7 @@ import torch.nn.functional as F
 import numpy as np
 from pathlib import Path
 from common.lut import forward_2x2_input_SxS_output
-from common.layers import RgbToYcbcr, YcbcrToRgb
+from common import layers 
 
 class SRLut(nn.Module):
     def __init__(
@@ -39,13 +39,13 @@ class SRLut(nn.Module):
 
 
 
-class SRLutRot90(nn.Module):
+class SRLutR90(nn.Module):
     def __init__(
         self, 
         quantization_interval,
         scale
     ):
-        super(SRLutRot90, self).__init__()     
+        super(SRLutR90, self).__init__()     
         self.scale = scale
         self.quantization_interval = quantization_interval
         self.stage_lut = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
@@ -77,18 +77,18 @@ class SRLutRot90(nn.Module):
         return f"{self.__class__.__name__}\n  lut size: {self.stage_lut.shape}"
 
 
-class SRLutRot90Y(nn.Module):
+class SRLutR90Y(nn.Module):
     def __init__(
         self, 
         quantization_interval,
         scale
     ):
-        super(SRLutRot90Y, self).__init__()     
+        super(SRLutR90Y, self).__init__()     
         self.scale = scale
         self.quantization_interval = quantization_interval
         self.stage_lut = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
-        self.rgb_to_ycbcr = RgbToYcbcr()
-        self.ycbcr_to_rgb = YcbcrToRgb()
+        self.rgb_to_ycbcr = layers.RgbToYcbcr()
+        self.ycbcr_to_rgb = layers.YcbcrToRgb()
 
     @staticmethod
     def init_from_lut(