The rest of the model (it didn’t fit in my last reply), and the memory summary:
    (deep_attention): MCA_ED(
      (enc_list): ModuleList(
        (0): SA(
          (mhatt): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (ffn): FFN(
            (mlp): MLP(
              (fc): FC(
                (linear): Linear(in_features=256, out_features=512, bias=True)
                (relu): ReLU(inplace=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (linear): Linear(in_features=512, out_features=256, bias=True)
            )
          )
          (dropout1): Dropout(p=0.1, inplace=False)
          (norm1): LayerNorm()
          (dropout2): Dropout(p=0.1, inplace=False)
          (norm2): LayerNorm()
        )
        (1): SA(
          (mhatt): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (ffn): FFN(
            (mlp): MLP(
              (fc): FC(
                (linear): Linear(in_features=256, out_features=512, bias=True)
                (relu): ReLU(inplace=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (linear): Linear(in_features=512, out_features=256, bias=True)
            )
          )
          (dropout1): Dropout(p=0.1, inplace=False)
          (norm1): LayerNorm()
          (dropout2): Dropout(p=0.1, inplace=False)
          (norm2): LayerNorm()
        )
        (2): SA(
          (mhatt): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (ffn): FFN(
            (mlp): MLP(
              (fc): FC(
                (linear): Linear(in_features=256, out_features=512, bias=True)
                (relu): ReLU(inplace=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (linear): Linear(in_features=512, out_features=256, bias=True)
            )
          )
          (dropout1): Dropout(p=0.1, inplace=False)
          (norm1): LayerNorm()
          (dropout2): Dropout(p=0.1, inplace=False)
          (norm2): LayerNorm()
        )
        (3): SA(
          (mhatt): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (ffn): FFN(
            (mlp): MLP(
              (fc): FC(
                (linear): Linear(in_features=256, out_features=512, bias=True)
                (relu): ReLU(inplace=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (linear): Linear(in_features=512, out_features=256, bias=True)
            )
          )
          (dropout1): Dropout(p=0.1, inplace=False)
          (norm1): LayerNorm()
          (dropout2): Dropout(p=0.1, inplace=False)
          (norm2): LayerNorm()
        )
        (4): SA(
          (mhatt): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (ffn): FFN(
            (mlp): MLP(
              (fc): FC(
                (linear): Linear(in_features=256, out_features=512, bias=True)
                (relu): ReLU(inplace=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (linear): Linear(in_features=512, out_features=256, bias=True)
            )
          )
          (dropout1): Dropout(p=0.1, inplace=False)
          (norm1): LayerNorm()
          (dropout2): Dropout(p=0.1, inplace=False)
          (norm2): LayerNorm()
        )
        (5): SA(
          (mhatt): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (ffn): FFN(
            (mlp): MLP(
              (fc): FC(
                (linear): Linear(in_features=256, out_features=512, bias=True)
                (relu): ReLU(inplace=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (linear): Linear(in_features=512, out_features=256, bias=True)
            )
          )
          (dropout1): Dropout(p=0.1, inplace=False)
          (norm1): LayerNorm()
          (dropout2): Dropout(p=0.1, inplace=False)
          (norm2): LayerNorm()
        )
      )
      (dec_list): ModuleList(
        (0): SGA(
          (mhatt1): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (mhatt2): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (ffn): FFN(
            (mlp): MLP(
              (fc): FC(
                (linear): Linear(in_features=256, out_features=512, bias=True)
                (relu): ReLU(inplace=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (linear): Linear(in_features=512, out_features=256, bias=True)
            )
          )
          (dropout1): Dropout(p=0.1, inplace=False)
          (norm1): LayerNorm()
          (dropout2): Dropout(p=0.1, inplace=False)
          (norm2): LayerNorm()
          (dropout3): Dropout(p=0.1, inplace=False)
          (norm3): LayerNorm()
        )
        (1): SGA(
          (mhatt1): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (mhatt2): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (ffn): FFN(
            (mlp): MLP(
              (fc): FC(
                (linear): Linear(in_features=256, out_features=512, bias=True)
                (relu): ReLU(inplace=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (linear): Linear(in_features=512, out_features=256, bias=True)
            )
          )
          (dropout1): Dropout(p=0.1, inplace=False)
          (norm1): LayerNorm()
          (dropout2): Dropout(p=0.1, inplace=False)
          (norm2): LayerNorm()
          (dropout3): Dropout(p=0.1, inplace=False)
          (norm3): LayerNorm()
        )
        (2): SGA(
          (mhatt1): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (mhatt2): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (ffn): FFN(
            (mlp): MLP(
              (fc): FC(
                (linear): Linear(in_features=256, out_features=512, bias=True)
                (relu): ReLU(inplace=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (linear): Linear(in_features=512, out_features=256, bias=True)
            )
          )
          (dropout1): Dropout(p=0.1, inplace=False)
          (norm1): LayerNorm()
          (dropout2): Dropout(p=0.1, inplace=False)
          (norm2): LayerNorm()
          (dropout3): Dropout(p=0.1, inplace=False)
          (norm3): LayerNorm()
        )
        (3): SGA(
          (mhatt1): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (mhatt2): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (ffn): FFN(
            (mlp): MLP(
              (fc): FC(
                (linear): Linear(in_features=256, out_features=512, bias=True)
                (relu): ReLU(inplace=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (linear): Linear(in_features=512, out_features=256, bias=True)
            )
          )
          (dropout1): Dropout(p=0.1, inplace=False)
          (norm1): LayerNorm()
          (dropout2): Dropout(p=0.1, inplace=False)
          (norm2): LayerNorm()
          (dropout3): Dropout(p=0.1, inplace=False)
          (norm3): LayerNorm()
        )
        (4): SGA(
          (mhatt1): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (mhatt2): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (ffn): FFN(
            (mlp): MLP(
              (fc): FC(
                (linear): Linear(in_features=256, out_features=512, bias=True)
                (relu): ReLU(inplace=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (linear): Linear(in_features=512, out_features=256, bias=True)
            )
          )
          (dropout1): Dropout(p=0.1, inplace=False)
          (norm1): LayerNorm()
          (dropout2): Dropout(p=0.1, inplace=False)
          (norm2): LayerNorm()
          (dropout3): Dropout(p=0.1, inplace=False)
          (norm3): LayerNorm()
        )
        (5): SGA(
          (mhatt1): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (mhatt2): MHAtt(
            (linear_v): Linear(in_features=256, out_features=256, bias=True)
            (linear_k): Linear(in_features=256, out_features=256, bias=True)
            (linear_q): Linear(in_features=256, out_features=256, bias=True)
            (linear_merge): Linear(in_features=256, out_features=256, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (ffn): FFN(
            (mlp): MLP(
              (fc): FC(
                (linear): Linear(in_features=256, out_features=512, bias=True)
                (relu): ReLU(inplace=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (linear): Linear(in_features=512, out_features=256, bias=True)
            )
          )
          (dropout1): Dropout(p=0.1, inplace=False)
          (norm1): LayerNorm()
          (dropout2): Dropout(p=0.1, inplace=False)
          (norm2): LayerNorm()
          (dropout3): Dropout(p=0.1, inplace=False)
          (norm3): LayerNorm()
        )
      )
    )
    (embedding_expansion): Linear(in_features=150, out_features=256, bias=True)
    (embedding_reduction): Linear(in_features=768, out_features=256, bias=True)
    (x_reshape): Sequential(
      (0): Flatten()
      (1): Linear(in_features=2560, out_features=1024, bias=True)
    )
    (y_reshape): Sequential(
      (0): Flatten()
      (1): Linear(in_features=246016, out_features=1024, bias=True)
    )
  )

|===========================================================================|
|                  PyTorch CUDA memory summary, device ID 0                 |
|---------------------------------------------------------------------------|
|            CUDA OOMs: 0            |        cudaMalloc retries: 0         |
|===========================================================================|
|        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |
|---------------------------------------------------------------------------|
| Allocated memory      |    3251 MB |    5549 MB |  344050 MB |  340799 MB |
|       from large pool |    3134 MB |    5387 MB |  333656 MB |  330522 MB |
|       from small pool |     117 MB |     167 MB |   10394 MB |   10277 MB |
|---------------------------------------------------------------------------|
| Active memory         |    3251 MB |    5549 MB |  344050 MB |  340799 MB |
|       from large pool |    3134 MB |    5387 MB |  333656 MB |  330522 MB |
|       from small pool |     117 MB |     167 MB |   10394 MB |   10277 MB |
|---------------------------------------------------------------------------|
| GPU reserved memory   |    6350 MB |    6350 MB |    6350 MB |       0 B  |
|       from large pool |    6182 MB |    6182 MB |    6182 MB |       0 B  |
|       from small pool |     168 MB |     168 MB |     168 MB |       0 B  |
|---------------------------------------------------------------------------|
| Non-releasable memory |     960 MB |    1056 MB |  259177 MB |  258216 MB |
|       from large pool |     941 MB |    1035 MB |  248215 MB |  247273 MB |
|       from small pool |      18 MB |      30 MB |   10961 MB |   10942 MB |
|---------------------------------------------------------------------------|
| Allocations           |    1330    |    2045    |  167490    |  166160    |
|       from large pool |      50    |     267    |   49870    |   49820    |
|       from small pool |    1280    |    1779    |  117620    |  116340    |
|---------------------------------------------------------------------------|
| Active allocs         |    1330    |    2045    |  167490    |  166160    |
|       from large pool |      50    |     267    |   49870    |   49820    |
|       from small pool |    1280    |    1779    |  117620    |  116340    |
|---------------------------------------------------------------------------|
| GPU reserved segments |     107    |     107    |     107    |       0    |
|       from large pool |      23    |      23    |      23    |       0    |
|       from small pool |      84    |      84    |      84    |       0    |
|---------------------------------------------------------------------------|
| Non-releasable allocs |     108    |     122    |   80164    |   80056    |
|       from large pool |      10    |      25    |   26731    |   26721    |
|       from small pool |      98    |     111    |   53433    |   53335    |
|===========================================================================|