Function clones

src/shape_assembly/models/train/transformer.py:8–9

(module, N)

import pytorch_lightning as pl

def clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

def attention(query, key, value, mask=None):
    d_k = query.size(-1)