用pytorch写一段四元数卷积

下面是一段用 PyTorch 实现四元数卷积的代码示例：

import torch
import torch.nn as nn
import torch.nn.functional as F

class QuaternionConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, bias=True):
        super(QuaternionConv2d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.bias = bias
        
        # 四元数卷积的权值是复数，因此需要 4 倍输入通道数
        self.weight = nn.Parameter(torch.Tensor(out_channels, in_channels * 4, kernel_size, kernel_size))
        if bias:
            self.bias = nn.Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()
        
    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in)
            nn.init.uniform_(self.bias, -bound, bound)
            
    def forward(self, input):
        # 将四元数展开为 4 个通道
        input = input.reshape(input.size(0), self.in_channels, 4, input.size(2), input.size(3))
        real, imag = input[:, :, 0], input[:, :, 1:]
        real, imag = real.permute(0, 2, 1, 3, 4), imag.permute(0, 2, 1, 3, 4)
        
        # 实数部分与虚数部分分别卷积
        real_out = F.conv2d(real, self.weight[:, :self.in_channels, :, :], stride=self.stride, padding=self.padding)
        imag_out = F.conv2d(imag, self.weight[:, self.in_channels:, :, :], stride=self.stride, padding=self.padding)
        
        # 将实数部分与虚数部分合并为四元数
        out = torch.stack([real_out, imag_out[