是否可以在Pytorch nn.Sequential()中添加条件

import torch
import torch.nn as nn

class CustomModel(nn.Module):
    def __init__(self):
        super(CustomModel, self).__init__()
        self.layer1 = nn.Linear(10, 20)
        self.layer2 = nn.Linear(20, 30)
        self.layer3 = nn.Linear(30, 2)

    def forward(self, x, condition):
        x = self.layer1(x)
        x = self.layer2(x)
        
        if condition:
            x = self.layer3(x)
        
        return x

# 创建模型实例
model = CustomModel()

# 使用模型
input_data = torch.randn(1, 10)
condition = True
output = model(input_data, condition)