| 804 | |
| 805 | def custom(start, end): |
| 806 | def custom_forward(*inputs): |
| 807 | layers_ = self.layers[start:end] |
| 808 | x_, inputs = inputs[0], inputs[1:] |
| 809 | if self.relative_encoding: |
| 810 | inputs, mems_ = inputs[:4], inputs[4:] |
| 811 | else: |
| 812 | inputs, mems_ = inputs[:1], inputs[1:] |
| 813 | for i, layer in enumerate(layers_): |
| 814 | mem_i_ = mems_[i] if mems_ else None |
| 815 | x_ = layer(x_, *inputs, mem=mem_i_) |
| 816 | if self.max_memory_length > 0 or return_memory: |
| 817 | mem_layers.append(check_detach(x_)) |
| 818 | return x_ |
| 819 | |
| 820 | return custom_forward |
| 821 | |