Solved by ariG23498 in post #2 
              
              
                Hey @DonghunP 
With the screenshot shared I am thinking the model should look something like this 
class CustomModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.projection = nn.Sequential(
            nn.Linear(64, 128),
            nn.Tanh(),
            nn.Linear(12…