Residual and Concatenate modules for PyTorch. Residual taken from https://github.com/lucidrains/denoising-diffusion-pytorch/blob/eb6e1b508e920efd840f2fa13e80d60ce19ed2c1/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py#L72 since it was very useful.

residual_concatenate.py

import torch
import torch.nn as nn

class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn
    def forward(self, x, *args, **kwargs):
        return self.fn(x, *args, **kwargs) + x

class Concat(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn
    def forward(self, x, *args, **kwargs):
        return torch.cat([x,self.fn(x, *args, **kwargs)],1)