(layer_num, memory_bank)
| 46 | return self_cache, mem_cache |
| 47 | |
| 48 | def init_onmt_cache(layer_num, memory_bank): |
| 49 | cache = {} |
| 50 | for i in range(layer_num): |
| 51 | layer_cache = {"memory_keys": None, "memory_values": None} |
| 52 | layer_cache["self_keys"] = None |
| 53 | layer_cache["self_values"] = None |
| 54 | cache[i] = layer_cache |
| 55 | return cache |
| 56 | |
| 57 | class ONMTDecoder(torch.nn.Module): |
| 58 | def __init__(self, layer_num, head_num, head_size, weights): |