Repository: pinder-org/pinder (language: python)

Issue:
Confusing features in dataset returned by `get_torch_loader`
I'm trying to use the dataloaders returned by `pinder.core.loader.dataset.get_torch_loader` for all-atom representations of the protein complexes. When I iterate over the keys present in a sample from one of these datasets, I get:
```
from pinder.core.loader.dataset import PinderDataset, get_torch_loader
train, val = [get_torch_loader(PinderDataset(split=split)) for split in ["train", "val"]]
for key in train.dataset[0].keys():
    print(key, type(train.dataset[0][key]))
    if type(train.dataset[0][key]) == dict:
        for key2 in train.dataset[0][key].keys():
            print("    ", key2, type(train.dataset[0][key][key2]), train.dataset[0][key][key2].shape)

### ------- Output ------- ###

target_complex <class 'dict'>
     atom_types <class 'torch.Tensor'> torch.Size([2556, 12])
     residue_types <class 'torch.Tensor'> torch.Size([2556, 1])
     atom_coordinates <class 'torch.Tensor'> torch.Size([2556, 3])
     residue_coordinates <class 'torch.Tensor'> torch.Size([2556, 3])
     residue_ids <class 'torch.Tensor'> torch.Size([2556])
     chain_ids <class 'torch.Tensor'> torch.Size([2556])
feature_complex <class 'dict'>
     atom_types <class 'torch.Tensor'> torch.Size([2556, 12])
     residue_types <class 'torch.Tensor'> torch.Size([2556, 1])
     atom_coordinates <class 'torch.Tensor'> torch.Size([2556, 3])
     residue_coordinates <class 'torch.Tensor'> torch.Size([2556, 3])
     residue_ids <class 'torch.Tensor'> torch.Size([2556])
     chain_ids <class 'torch.Tensor'> torch.Size([2556])
id <class 'str'>
sample_id <class 'str'>
target_id <class 'str'>
```
I have the following questions:

**What is the difference between `atom_coordinates` and `residue_coordinates`?**

For at least the first 10 examples in the train dataset, I find that
```
for i in range(10):
    print((train.dataset[i]['feature_complex']['residue_coordinates'] == train.dataset[i]['feature_complex']['atom_coordinates']).all())
```
returns `True`.

**Can the loader be modified to also return atom name information?**

The `atom_types` tensor provides only a one-hot encoding of element type, for example, running
```
train.dataset[0]['target_complex']['atom_types'].unique(dim=0)
```
returns
```
tensor([[0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
```
This is confusing when compared to the `Structure` class, in which `atom_name` is encoded as elements of `constants.BB_ATOMS` and `constants.SC_ATOMS` (e.g. 'N', 'CA', 'C', 'O', 'CB', 'CG', etc.).

**It would be useful/less confusing if the dataloader returned by `get_torch_loader` contained information that corresponded more closely to the data provided in the `Structure` class.**


Discussion / hints:
Hi @eswan01, thanks for the feedback on the example transform implementation. I agree that it would make sense to rename the misleading `atom_types` key to reflect that its actually encodings of elements and not atom types and also introduce an atom type encoding. 


## Regarding your first question

> What is the difference between atom_coordinates and residue_coordinates?

In the `PinderDataset` (torch dataset), there is no difference. It appears this was an oversight when hijacking the existing transform that was targeting the torch-geometric dataset. In the torch-geometric dataset, the difference is that `residue_coordinates` tracks only the C-alpha atom coordinates (one per residue) and `atom_coordinates` tracks all-atom. Note the `receptor.filter("atom_name", mask=["CA"])` below:
https://github.com/pinder-org/pinder/blob/bb62635bcb4186cf4f45696fc3da5082bb899f0a/src/pinder-core/pinder/core/loader/geodata.py#L118-L133

That `structure2tensor` function really just serves as an example transform for how one could go from `Structure` instances to Tensor-encoded data items. In the `PinderDataset` implementation, the _default_  torch-API transform defined is the `structure2tensor_transform` which essentially just wraps the `structure2tensor` method such that it accepts a `Structure` as input and returns a dict of tensors. 
https://github.com/pinder-org/pinder/blob/bb62635bcb4186cf4f45696fc3da5082bb899f0a/src/pinder-core/pinder/core/loader/dataset.py#L238-L254

As you can see below, the current version uses all-atom coordinates for both `residue_coordinates` and `atom_coordinates` (the oversight):
https://github.com/pinder-org/pinder/blob/bb62635bcb4186cf4f45696fc3da5082bb899f0a/src/pinder-core/pinder/core/loader/dataset.py#L54-L63

I will push a fix to that default transform shortly 👍 


## Regarding the second question

> Can the loader be modified to also return atom name information?

As I alluded to above, we can definitely adjust the info being returned by the default transform as you suggested. I will try to include that in the upcoming patch for the `residue_coordinates` bug. 

That being said, the short answer to 

> Can the loader be modified to also return atom name information? 
> It would be useful/less confusing if the dataloader returned by get_torch_loader contained information that corresponded more closely to the data provided in the Structure class.
> 

**Absolutely!** You can define whatever transform you desire, that's why the `PinderLoader` default return value of `__getitem__` contains a tuple of `(PinderSystem, Structure(<feature complex>), Structure(<target complex>))` 

https://github.com/pinder-org/pinder/blob/bb62635bcb4186cf4f45696fc3da5082bb899f0a/src/pinder-core/pinder/core/loader/dataset.py#L293-L299

You can specify your own `transform` and `target_transform` to apply to the `feature_complex` `Structure` and `target_complex` `Structure`. This closely follows the transform/target_transform concepts from the `torch` datasets + dataloader tutorial: https://pytorch.org/tutorials/beginner/basics/data_tutorial.html#creating-a-custom-dataset-for-your-files



Fix patch:
diff --git a/README.md b/README.md
index 66550ec..e8be078 100644
--- a/README.md
+++ b/README.md
@@ -375,7 +375,8 @@ Each model decoy should have exactly two chains: {R, L} for {Receptor, Ligand},
 }
 ```
 
-For more details on the implementations of the eval metrics, see the [eval docs](examples/eval/).
+For more details on the implementations of the eval metrics, see the [eval tutorial](https://pinder-org.github.io/pinder/pinder-eval.html), [API docs](https://pinder-org.github.io/pinder/source/pinder.eval.dockq.html#) and [eval FAQ](https://pinder-org.github.io/pinder/faq.html#how-can-i-use-the-evaluation-harness-outside-of-a-pinder-context).
+
 For more details on leaderboard generation, see the [MethodMetrics](src/pinder-eval/pinder/eval/dockq/method.py) implementation.
 
 
diff --git a/examples/pinder-loader.ipynb b/examples/pinder-loader.ipynb
index 47a563f..46af8a5 100644
--- a/examples/pinder-loader.ipynb
+++ b/examples/pinder-loader.ipynb
@@ -435,7 +435,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 7,
    "id": "6e09ccc9-215f-4d8e-a08f-1de96bb42131",
    "metadata": {},
    "outputs": [
@@ -445,7 +445,7 @@
        "PinderLoader(split=test, monomers=holo, systems=180)"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -481,7 +481,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 8,
    "id": "f7bcbe6e-4bd3-435c-9ff0-cf611f6bf9cf",
    "metadata": {},
    "outputs": [
@@ -491,7 +491,7 @@
        "180"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -502,7 +502,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 9,
    "id": "fecd34da-c91d-400c-bbda-64d6cccb8e0b",
    "metadata": {},
    "outputs": [
@@ -521,7 +521,7 @@
        " pinder.core.loader.structure.Structure)"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -535,7 +535,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 10,
    "id": "bdb3327d-26d2-46fc-95ea-de2ea2fc0516",
    "metadata": {},
    "outputs": [
@@ -543,7 +543,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 180/180 [01:14<00:00,  2.43it/s]\n"
+      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 180/180 [01:12<00:00,  2.50it/s]\n"
      ]
     }
    ],
@@ -565,7 +565,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 11,
    "id": "77d0f734-508a-4b1e-a4f3-acb8590869a0",
    "metadata": {},
    "outputs": [
@@ -575,7 +575,7 @@
        "set()"
       ]
      },
-     "execution_count": 10,
+     "execution_count": 11,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -601,7 +601,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 12,
    "id": "3644c869-da61-4620-a054-2a3ad163f3f8",
    "metadata": {},
    "outputs": [
@@ -611,7 +611,7 @@
        "True"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -637,7 +637,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 13,
    "id": "6b653c61-cebd-4ace-baf8-a27f1e011466",
    "metadata": {},
    "outputs": [
@@ -645,8 +645,8 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "[PosixPath('/var/folders/tt/x223wxwj6dzg3vjjgc_6y5bm0000gn/T/tmpbe5qjtfe/117e__A1_P00817--117e__B1_P00817/af__P00817.pdb')]\n",
-      "[PosixPath('/var/folders/tt/x223wxwj6dzg3vjjgc_6y5bm0000gn/T/tmpbe5qjtfe/1df0__A1_Q07009--1df0__B1_Q64537/af__Q07009.pdb'), PosixPath('/var/folders/tt/x223wxwj6dzg3vjjgc_6y5bm0000gn/T/tmpbe5qjtfe/1df0__A1_Q07009--1df0__B1_Q64537/af__Q64537.pdb')]\n"
+      "[PosixPath('/var/folders/tt/x223wxwj6dzg3vjjgc_6y5bm0000gn/T/tmpgrbzq_is/117e__A1_P00817--117e__B1_P00817/af__P00817.pdb')]\n",
+      "[PosixPath('/var/folders/tt/x223wxwj6dzg3vjjgc_6y5bm0000gn/T/tmpgrbzq_is/1df0__A1_Q07009--1df0__B1_Q64537/af__Q07009.pdb'), PosixPath('/var/folders/tt/x223wxwj6dzg3vjjgc_6y5bm0000gn/T/tmpgrbzq_is/1df0__A1_Q07009--1df0__B1_Q64537/af__Q64537.pdb')]\n"
      ]
     }
    ],
@@ -738,7 +738,8 @@
     "* `residue_types`\n",
     "* `residue_ids`\n",
     "\n",
-    "You can choose to use a different representation by overriding the default values of `transform` and `target_transform`.\n",
+    "**Note: You can choose to use a different representation by overriding the default values of `transform` and `target_transform`. The default transform is the structure2tensor_transform defined in pinder.core.loader.dataset. \n",
+    "It simply takes a `Structure` object and returns a dictionary with string keys and tensor values.**\n",
     " \n",
     "It leverages the `PinderLoader` to apply optional filters and/or transforms, provide an interface for sampling alternative monomers, and exposes `transform` and `target_transform` arguments used by the torch Dataset API. \n",
     "\n",
@@ -747,24 +748,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 14,
    "id": "9658c6aa-5720-4c0f-98ae-b79d34ad2754",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "<pinder.core.loader.dataset.PinderDataset at 0x1ecf92350>"
+       "<pinder.core.loader.dataset.PinderDataset at 0x1fc1f76a0>"
       ]
      },
-     "execution_count": 2,
+     "execution_count": 14,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "from pinder.core.loader import filters, transforms\n",
-    "from pinder.core.loader.dataset import PinderDataset\n",
+    "from pinder.core.loader.dataset import PinderDataset, structure2tensor_transform\n",
     "\n",
     "base_filters = [\n",
     "    filters.FilterByMissingHolo(),\n",
@@ -792,6 +793,9 @@
     "    structure_transforms_target=structure_transforms,\n",
     "    # Apply to the feature complex\n",
     "    structure_transforms_feature=structure_transforms,\n",
+    "    # This is the default transform if not specified\n",
+    "    transform=structure2tensor_transform,\n",
+    "    target_transform=structure2tensor_transform,\n",
     ")\n",
     "assert len(train_dataset) == len(get_index().query('split == \"train\"'))\n",
     "\n",
@@ -834,20 +838,27 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 15,
    "id": "383244d7-1a85-4732-9dd4-3455fe4cec88",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "{'target_complex': {'atom_types': tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
+       "{'target_complex': {'atom_types': tensor([[0.],\n",
+       "          [1.],\n",
+       "          [2.],\n",
+       "          ...,\n",
+       "          [1.],\n",
+       "          [2.],\n",
+       "          [3.]]),\n",
+       "  'element_types': tensor([[3.],\n",
+       "          [0.],\n",
+       "          [0.],\n",
        "          ...,\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [0., 0., 1.,  ..., 0., 0., 0.]]),\n",
+       "          [0.],\n",
+       "          [0.],\n",
+       "          [2.]]),\n",
        "  'residue_types': tensor([[16.],\n",
        "          [16.],\n",
        "          [16.],\n",
@@ -862,22 +873,351 @@
        "          [177.7620, 463.8650, 166.9020],\n",
        "          [177.4130, 465.0800, 167.7550],\n",
        "          [176.8000, 464.9490, 168.8150]]),\n",
-       "  'residue_coordinates': tensor([[131.7500, 429.3090, 163.5360],\n",
-       "          [132.6810, 428.2520, 163.1550],\n",
-       "          [133.5150, 428.6750, 161.9500],\n",
-       "          ...,\n",
-       "          [177.7620, 463.8650, 166.9020],\n",
-       "          [177.4130, 465.0800, 167.7550],\n",
-       "          [176.8000, 464.9490, 168.8150]]),\n",
+       "  'residue_coordinates': tensor([[132.6810, 428.2520, 163.1550],\n",
+       "          [133.5560, 429.2490, 159.5910],\n",
+       "          [133.8750, 432.8980, 160.6290],\n",
+       "          [136.1110, 431.8050, 163.5130],\n",
+       "          [138.3420, 429.8920, 161.0830],\n",
+       "          [138.5230, 432.9090, 158.7600],\n",
+       "          [139.4710, 435.1730, 161.6750],\n",
+       "          [142.1200, 432.6310, 162.6740],\n",
+       "          [143.5890, 432.7500, 159.1600],\n",
+       "          [143.6190, 436.5600, 159.1540],\n",
+       "          [145.2600, 436.8040, 162.5830],\n",
+       "          [147.8380, 434.1320, 161.7330],\n",
+       "          [148.7390, 435.8920, 158.4790],\n",
+       "          [149.0640, 439.2630, 160.2240],\n",
+       "          [151.2220, 437.7670, 162.9840],\n",
+       "          [153.4710, 435.8030, 160.6120],\n",
+       "          [153.9420, 438.9180, 158.4710],\n",
+       "          [156.1140, 440.2660, 161.3150],\n",
+       "          [158.3330, 437.1860, 161.4250],\n",
+       "          [161.8090, 436.4120, 160.1100],\n",
+       "          [161.1120, 432.8370, 158.9350],\n",
+       "          [157.4020, 432.0710, 159.3930],\n",
+       "          [156.5290, 428.3870, 159.0310],\n",
+       "          [154.0400, 428.1160, 156.1500],\n",
+       "          [153.3520, 424.6590, 154.7510],\n",
+       "          [152.8370, 424.8210, 150.9870],\n",
+       "          [150.9700, 422.4160, 148.7320],\n",
+       "          [153.5900, 422.6860, 145.9670],\n",
+       "          [157.3050, 423.3410, 145.6600],\n",
+       "          [158.4860, 426.6730, 144.2700],\n",
+       "          [160.7730, 427.0880, 141.2590],\n",
+       "          [163.9630, 429.0530, 141.9600],\n",
+       "          [166.4230, 427.4330, 139.5550],\n",
+       "          [167.6170, 430.8330, 138.3250],\n",
+       "          [170.5180, 431.9210, 140.5320],\n",
+       "          [169.7800, 435.6680, 140.3160],\n",
+       "          [167.6790, 437.8790, 142.6110],\n",
+       "          [167.3350, 435.1910, 145.2890],\n",
+       "          [168.9090, 437.3000, 148.0580],\n",
+       "          [166.1860, 439.9440, 148.5750],\n",
+       "          [164.1010, 438.1300, 151.1890],\n",
+       "          [163.0190, 438.8690, 154.7450],\n",
+       "          [164.6750, 435.6560, 155.9860],\n",
+       "          [168.1170, 436.6550, 154.6330],\n",
+       "          [169.8040, 437.0930, 158.0010],\n",
+       "          [173.2680, 438.4850, 158.5750],\n",
+       "          [174.9270, 440.8490, 156.1460],\n",
+       "          [175.9170, 440.8010, 152.5000],\n",
+       "          [179.6060, 441.0910, 153.4330],\n",
+       "          [180.8000, 439.9830, 156.8710],\n",
+       "          [184.2540, 439.9980, 158.4580],\n",
+       "          [185.5840, 437.2970, 160.7670],\n",
+       "          [186.4140, 439.9000, 163.4480],\n",
+       "          [182.9430, 441.4570, 163.7500],\n",
+       "          [181.4930, 442.0910, 167.2100],\n",
+       "          [178.0780, 440.5360, 167.9060],\n",
+       "          [176.1650, 440.9290, 171.1790],\n",
+       "          [173.0530, 438.9530, 172.1310],\n",
+       "          [170.3990, 440.6990, 174.2170],\n",
+       "          [166.7880, 440.0710, 175.2260],\n",
+       "          [164.4480, 442.1010, 173.0340],\n",
+       "          [160.9780, 443.6110, 173.3810],\n",
+       "          [158.8330, 443.9000, 170.2730],\n",
+       "          [160.1820, 442.9360, 166.8360],\n",
+       "          [163.5010, 444.7230, 166.1630],\n",
+       "          [164.6310, 442.4290, 163.3160],\n",
+       "          [166.6330, 444.2560, 160.6240],\n",
+       "          [166.1700, 447.5210, 162.5220],\n",
+       "          [168.1350, 449.9450, 164.6440],\n",
+       "          [168.0040, 449.0470, 168.3240],\n",
+       "          [168.4210, 451.0540, 171.5220],\n",
+       "          [169.0440, 450.0360, 175.1300],\n",
+       "          [166.1650, 450.3680, 177.6100],\n",
+       "          [166.7800, 449.8300, 181.3340],\n",
+       "          [163.8480, 448.9150, 183.5760],\n",
+       "          [163.3500, 443.9770, 187.9750],\n",
+       "          [165.0240, 442.4930, 184.8870],\n",
+       "          [168.1120, 443.0980, 182.7810],\n",
+       "          [168.1280, 445.9210, 180.2050],\n",
+       "          [166.7230, 445.0130, 176.8030],\n",
+       "          [167.0690, 446.1330, 173.1930],\n",
+       "          [164.0790, 447.8390, 171.5780],\n",
+       "          [163.4630, 449.0630, 168.0350],\n",
+       "          [164.4160, 452.7320, 168.0130],\n",
+       "          [166.7460, 455.3190, 166.5790],\n",
+       "          [167.5460, 459.0090, 166.8440],\n",
+       "          [170.2010, 459.9620, 169.3810],\n",
+       "          [169.5150, 456.8220, 171.4570],\n",
+       "          [170.7440, 454.3880, 168.7860],\n",
+       "          [172.8900, 451.6810, 170.3850],\n",
+       "          [173.2420, 449.3000, 167.4490],\n",
+       "          [171.4700, 447.3670, 164.7150],\n",
+       "          [169.7840, 443.9650, 164.9670],\n",
+       "          [170.8880, 441.4120, 162.3680],\n",
+       "          [169.3320, 438.1640, 163.6510],\n",
+       "          [166.7170, 436.9060, 166.1030],\n",
+       "          [166.2830, 433.5290, 167.7900],\n",
+       "          [162.5490, 433.4260, 168.5040],\n",
+       "          [162.6440, 430.3540, 170.7620],\n",
+       "          [164.9490, 432.1120, 173.2350],\n",
+       "          [163.6530, 435.5760, 172.2210],\n",
+       "          [167.2540, 436.7450, 171.8300],\n",
+       "          [168.3710, 439.3650, 169.3060],\n",
+       "          [171.8910, 439.6680, 167.8930],\n",
+       "          [173.1360, 443.2500, 167.5870],\n",
+       "          [176.0990, 444.7950, 165.8030],\n",
+       "          [177.1670, 447.8660, 167.8270],\n",
+       "          [177.4160, 451.2430, 166.1360],\n",
+       "          [181.1790, 451.0810, 166.7480],\n",
+       "          [181.3560, 448.4350, 163.9980],\n",
+       "          [180.6020, 448.8920, 160.3060],\n",
+       "          [177.8440, 446.8990, 158.6170],\n",
+       "          [176.7540, 446.3160, 155.0150],\n",
+       "          [173.3820, 444.6900, 154.4770],\n",
+       "          [169.6430, 444.9660, 153.9760],\n",
+       "          [167.6080, 446.9450, 156.5140],\n",
+       "          [163.9630, 447.9270, 156.8800],\n",
+       "          [163.1520, 451.2160, 155.1430],\n",
+       "          [160.3870, 453.6570, 156.0420],\n",
+       "          [158.9330, 453.6920, 152.5150],\n",
+       "          [159.8050, 453.2460, 148.8360],\n",
+       "          [161.2560, 456.7540, 148.4440],\n",
+       "          [164.9420, 455.8330, 148.8520],\n",
+       "          [167.3780, 455.9500, 145.9260],\n",
+       "          [171.0380, 455.0610, 145.5060],\n",
+       "          [173.5830, 457.5340, 146.9270],\n",
+       "          [171.2360, 458.9430, 149.5740],\n",
+       "          [171.8150, 459.8640, 153.2130],\n",
+       "          [169.5910, 457.9090, 155.6000],\n",
+       "          [168.4960, 458.3960, 159.2100],\n",
+       "          [167.0870, 455.9020, 161.7110],\n",
+       "          [163.6430, 457.0390, 162.8610],\n",
+       "          [161.8700, 456.3550, 166.1660],\n",
+       "          [160.7490, 452.9260, 164.8950],\n",
+       "          [164.2440, 451.8750, 163.8110],\n",
+       "          [163.4490, 452.3040, 160.1090],\n",
+       "          [165.6820, 454.0140, 157.5590],\n",
+       "          [164.2940, 457.2310, 156.0670],\n",
+       "          [165.6810, 459.7290, 153.5780],\n",
+       "          [167.3870, 462.7790, 155.0760],\n",
+       "          [170.2060, 465.8340, 161.8360],\n",
+       "          [172.1620, 462.7680, 160.6890],\n",
+       "          [174.0280, 461.7040, 163.8180],\n",
+       "          [175.2110, 458.5660, 161.9930],\n",
+       "          [176.5410, 458.2410, 158.4450],\n",
+       "          [174.5550, 455.7530, 156.3540],\n",
+       "          [174.3580, 455.6030, 152.5600],\n",
+       "          [171.8590, 453.6260, 150.4910],\n",
+       "          [173.6300, 451.2730, 148.0810],\n",
+       "          [170.4820, 450.1750, 146.2000],\n",
+       "          [166.8520, 451.1980, 145.7810],\n",
+       "          [164.0730, 450.1880, 148.1490],\n",
+       "          [162.6710, 446.7750, 147.1980],\n",
+       "          [159.1420, 445.8080, 148.2240],\n",
+       "          [158.3160, 442.3640, 149.5950],\n",
+       "          [155.0300, 443.1470, 151.4050],\n",
+       "          [152.6800, 446.0960, 151.8500],\n",
+       "          [155.0610, 447.4890, 154.4910],\n",
+       "          [158.2010, 445.3460, 154.0190],\n",
+       "          [160.8090, 447.5340, 152.2870],\n",
+       "          [164.4130, 446.3150, 152.1360],\n",
+       "          [167.3360, 448.5380, 151.1380],\n",
+       "          [171.0310, 447.6630, 150.9500],\n",
+       "          [172.8160, 450.1440, 153.2260],\n",
+       "          [176.3550, 450.4540, 154.5850],\n",
+       "          [176.6310, 452.0780, 158.0240],\n",
+       "          [180.1600, 453.0970, 159.0310],\n",
+       "          [179.9110, 454.9710, 162.3160],\n",
+       "          [178.7370, 458.4030, 163.4610],\n",
+       "          [179.1350, 461.9150, 162.0830],\n",
+       "          [180.3160, 464.9940, 163.9950],\n",
+       "          [195.1700, 454.0210, 194.1160],\n",
+       "          [196.4220, 450.5570, 193.1610],\n",
+       "          [195.0530, 447.5260, 191.3340],\n",
+       "          [195.7530, 445.4320, 194.4410],\n",
+       "          [195.6760, 446.8960, 197.9710],\n",
+       "          [199.1230, 447.9300, 199.1640],\n",
+       "          [198.4070, 447.0140, 202.7840],\n",
+       "          [198.1130, 443.3080, 203.5090],\n",
+       "          [195.0350, 441.7530, 205.1060],\n",
+       "          [194.7040, 440.1720, 208.5300],\n",
+       "          [191.7420, 437.9460, 207.6670],\n",
+       "          [190.9460, 435.0990, 205.3140],\n",
+       "          [188.0870, 435.6760, 202.8880],\n",
+       "          [184.8300, 433.7320, 203.1830],\n",
+       "          [184.1270, 431.8260, 199.9570],\n",
+       "          [180.5450, 430.8500, 200.7460],\n",
+       "          [178.2470, 433.3320, 199.0220],\n",
+       "          [178.2770, 433.9260, 195.2680],\n",
+       "          [176.3170, 437.2030, 195.0280],\n",
+       "          [177.7690, 440.7330, 194.9660],\n",
+       "          [181.3350, 439.3910, 195.1330],\n",
+       "          [182.5820, 441.3460, 192.0940],\n",
+       "          [182.3530, 444.9560, 193.3610],\n",
+       "          [186.0870, 445.2240, 194.0290],\n",
+       "          [188.7230, 447.5570, 192.6190],\n",
+       "          [190.8760, 444.5870, 191.5490],\n",
+       "          [188.2670, 443.5230, 188.9590],\n",
+       "          [190.1360, 444.3010, 185.7360],\n",
+       "          [188.8140, 444.1990, 182.1980],\n",
+       "          [185.1300, 444.1560, 181.3470],\n",
+       "          [182.1290, 442.2030, 182.6060],\n",
+       "          [181.3330, 441.2700, 178.9860],\n",
+       "          [184.3690, 440.6570, 176.7700],\n",
+       "          [184.4240, 439.7520, 173.0740],\n",
+       "          [187.1540, 437.5560, 171.6020],\n",
+       "          [187.5170, 439.9750, 168.6500],\n",
+       "          [188.4180, 443.0500, 170.7240],\n",
+       "          [191.4360, 445.2100, 169.8910],\n",
+       "          [193.7480, 445.6510, 172.8940],\n",
+       "          [196.8200, 447.8990, 173.0280],\n",
+       "          [199.4910, 448.6120, 175.6310],\n",
+       "          [201.2080, 451.9090, 176.3780],\n",
+       "          [203.6110, 453.4740, 178.8680],\n",
+       "          [201.6130, 454.9210, 181.7570],\n",
+       "          [202.0690, 458.3840, 183.2730],\n",
+       "          [200.4970, 458.3690, 186.7180],\n",
+       "          [197.7960, 455.8130, 187.5240],\n",
+       "          [195.1460, 455.8470, 184.7720],\n",
+       "          [193.6160, 452.5740, 185.9790],\n",
+       "          [189.8600, 452.3800, 185.3150],\n",
+       "          [189.9440, 455.8580, 183.7680],\n",
+       "          [189.7400, 457.4310, 180.3410],\n",
+       "          [193.1630, 457.9480, 178.8100],\n",
+       "          [194.6770, 460.5170, 176.4640],\n",
+       "          [197.9720, 460.6430, 174.6010],\n",
+       "          [200.7710, 462.5710, 176.3210],\n",
+       "          [203.7970, 463.7400, 174.3330],\n",
+       "          [207.1150, 464.5820, 175.9770],\n",
+       "          [210.1390, 466.5260, 174.8060],\n",
+       "          [210.8850, 463.9390, 172.1200],\n",
+       "          [213.6440, 461.8050, 173.6150],\n",
+       "          [211.1280, 459.2430, 174.9140],\n",
+       "          [208.0390, 457.7610, 173.2930],\n",
+       "          [204.4770, 458.9110, 173.9310],\n",
+       "          [202.5240, 457.8000, 177.0030],\n",
+       "          [198.9190, 457.7220, 178.2440],\n",
+       "          [197.5730, 459.9040, 181.0500],\n",
+       "          [194.2310, 460.2350, 182.8060],\n",
+       "          [192.1140, 462.6430, 180.7850],\n",
+       "          [189.0250, 463.1720, 178.7000],\n",
+       "          [187.1240, 465.8360, 176.8220],\n",
+       "          [188.2850, 466.3320, 173.2420],\n",
+       "          [191.5580, 464.4520, 173.8600],\n",
+       "          [189.9950, 461.1280, 174.9010],\n",
+       "          [191.8440, 458.2970, 173.1540],\n",
+       "          [190.4670, 455.2690, 174.9900],\n",
+       "          [189.7380, 453.5820, 178.3070],\n",
+       "          [192.3390, 451.8830, 180.5000],\n",
+       "          [191.2020, 448.5340, 181.9050],\n",
+       "          [194.3800, 447.0280, 183.4050],\n",
+       "          [197.8700, 447.9650, 184.5850],\n",
+       "          [200.9840, 445.8270, 184.9660],\n",
+       "          [202.8290, 447.6220, 187.7590],\n",
+       "          [206.2120, 445.9180, 187.3400],\n",
+       "          [206.6040, 447.1310, 183.7460],\n",
+       "          [204.4430, 450.2380, 184.1830],\n",
+       "          [202.2310, 449.1370, 181.2890],\n",
+       "          [198.6660, 450.3790, 180.7720],\n",
+       "          [196.2310, 448.2490, 178.7680],\n",
+       "          [194.0550, 450.4240, 176.5260],\n",
+       "          [190.8490, 449.7090, 174.6070],\n",
+       "          [190.4360, 452.3080, 171.8240],\n",
+       "          [187.3830, 454.5590, 171.7980],\n",
+       "          [186.4670, 452.9730, 168.4480],\n",
+       "          [185.9050, 449.6480, 170.2580],\n",
+       "          [182.8890, 448.8910, 172.4330],\n",
+       "          [183.5680, 448.6030, 176.1630],\n",
+       "          [181.3330, 447.5400, 179.0590],\n",
+       "          [182.7340, 447.5270, 182.5670],\n",
+       "          [183.3280, 449.2260, 185.8920],\n",
+       "          [184.8760, 452.6980, 185.6990],\n",
+       "          [185.7140, 455.4170, 188.2030],\n",
+       "          [182.9420, 457.9530, 188.8300],\n",
+       "          [183.2930, 461.5650, 189.9350],\n",
+       "          [182.7580, 462.1320, 193.6660],\n",
+       "          [180.0620, 464.7490, 193.1030],\n",
+       "          [176.4120, 463.9030, 192.5790],\n",
+       "          [177.1550, 460.1730, 192.6930],\n",
+       "          [173.8180, 459.5980, 194.4390],\n",
+       "          [172.0770, 461.1030, 191.4040],\n",
+       "          [173.6950, 458.5410, 189.0910],\n",
+       "          [171.0060, 456.0190, 188.1270],\n",
+       "          [170.8210, 453.3390, 185.4210],\n",
+       "          [170.0180, 454.8270, 182.0350],\n",
+       "          [171.8690, 458.1080, 182.6590],\n",
+       "          [174.2320, 459.6520, 180.1210],\n",
+       "          [177.8640, 459.9130, 181.2230],\n",
+       "          [180.8780, 461.8410, 179.9490],\n",
+       "          [184.5650, 461.8070, 180.8430],\n",
+       "          [186.0650, 464.7990, 182.6440],\n",
+       "          [189.5780, 466.2740, 182.6170],\n",
+       "          [190.7120, 463.6340, 185.1390],\n",
+       "          [189.2980, 460.6630, 183.2290],\n",
+       "          [186.3560, 460.3000, 185.6240],\n",
+       "          [182.7840, 459.6590, 184.5210],\n",
+       "          [180.2130, 462.3320, 185.3590],\n",
+       "          [176.5320, 462.8410, 184.6140],\n",
+       "          [175.8580, 465.0300, 181.5790],\n",
+       "          [178.1930, 465.7110, 174.9000],\n",
+       "          [180.9150, 464.1190, 172.7670],\n",
+       "          [181.2510, 460.5420, 174.1130],\n",
+       "          [178.4430, 457.9660, 174.2140],\n",
+       "          [178.5420, 456.4790, 177.7180],\n",
+       "          [175.4730, 455.1590, 179.5500],\n",
+       "          [175.4670, 454.0690, 183.1870],\n",
+       "          [174.1210, 450.5430, 183.6830],\n",
+       "          [173.9850, 450.6340, 187.5030],\n",
+       "          [174.3640, 452.9240, 190.5240],\n",
+       "          [177.5140, 454.2200, 192.1820],\n",
+       "          [179.3060, 451.6570, 194.3540],\n",
+       "          [181.4890, 453.2110, 197.0570],\n",
+       "          [184.7390, 451.3110, 197.6330],\n",
+       "          [186.4870, 454.0230, 199.6690],\n",
+       "          [185.5270, 457.3810, 201.1790],\n",
+       "          [186.7910, 459.0050, 197.9610],\n",
+       "          [186.6010, 456.0430, 195.5320],\n",
+       "          [183.3680, 455.3620, 193.6270],\n",
+       "          [182.9110, 452.9430, 190.7280],\n",
+       "          [180.0190, 452.5230, 188.3060],\n",
+       "          [179.0910, 450.0250, 185.5960],\n",
+       "          [178.9850, 451.7840, 182.2280],\n",
+       "          [179.0510, 451.0080, 178.5100],\n",
+       "          [181.0530, 453.1550, 176.0780],\n",
+       "          [180.1320, 452.7170, 172.4060],\n",
+       "          [182.0570, 455.4530, 170.6120],\n",
+       "          [181.8140, 459.1820, 170.0770],\n",
+       "          [178.5990, 461.0580, 169.3130],\n",
+       "          [177.7620, 463.8650, 166.9020]]),\n",
        "  'residue_ids': tensor([  4.,   4.,   4.,  ..., 182., 182., 182.]),\n",
        "  'chain_ids': tensor([0., 0., 0.,  ..., 1., 1., 1.])},\n",
-       " 'feature_complex': {'atom_types': tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
+       " 'feature_complex': {'atom_types': tensor([[0.],\n",
+       "          [1.],\n",
+       "          [2.],\n",
+       "          ...,\n",
+       "          [1.],\n",
+       "          [2.],\n",
+       "          [3.]]),\n",
+       "  'element_types': tensor([[3.],\n",
+       "          [0.],\n",
+       "          [0.],\n",
        "          ...,\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [0., 0., 1.,  ..., 0., 0., 0.]]),\n",
+       "          [0.],\n",
+       "          [0.],\n",
+       "          [2.]]),\n",
        "  'residue_types': tensor([[16.],\n",
        "          [16.],\n",
        "          [16.],\n",
@@ -892,13 +1232,335 @@
        "          [177.7620, 463.8650, 166.9020],\n",
        "          [177.4130, 465.0800, 167.7550],\n",
        "          [176.8000, 464.9490, 168.8150]]),\n",
-       "  'residue_coordinates': tensor([[131.7500, 429.3090, 163.5360],\n",
-       "          [132.6810, 428.2520, 163.1550],\n",
-       "          [133.5150, 428.6750, 161.9500],\n",
-       "          ...,\n",
-       "          [177.7620, 463.8650, 166.9020],\n",
-       "          [177.4130, 465.0800, 167.7550],\n",
-       "          [176.8000, 464.9490, 168.8150]]),\n",
+       "  'residue_coordinates': tensor([[132.6810, 428.2520, 163.1550],\n",
+       "          [133.5560, 429.2490, 159.5910],\n",
+       "          [133.8750, 432.8980, 160.6290],\n",
+       "          [136.1110, 431.8050, 163.5130],\n",
+       "          [138.3420, 429.8920, 161.0830],\n",
+       "          [138.5230, 432.9090, 158.7600],\n",
+       "          [139.4710, 435.1730, 161.6750],\n",
+       "          [142.1200, 432.6310, 162.6740],\n",
+       "          [143.5890, 432.7500, 159.1600],\n",
+       "          [143.6190, 436.5600, 159.1540],\n",
+       "          [145.2600, 436.8040, 162.5830],\n",
+       "          [147.8380, 434.1320, 161.7330],\n",
+       "          [148.7390, 435.8920, 158.4790],\n",
+       "          [149.0640, 439.2630, 160.2240],\n",
+       "          [151.2220, 437.7670, 162.9840],\n",
+       "          [153.4710, 435.8030, 160.6120],\n",
+       "          [153.9420, 438.9180, 158.4710],\n",
+       "          [156.1140, 440.2660, 161.3150],\n",
+       "          [158.3330, 437.1860, 161.4250],\n",
+       "          [161.8090, 436.4120, 160.1100],\n",
+       "          [161.1120, 432.8370, 158.9350],\n",
+       "          [157.4020, 432.0710, 159.3930],\n",
+       "          [156.5290, 428.3870, 159.0310],\n",
+       "          [154.0400, 428.1160, 156.1500],\n",
+       "          [153.3520, 424.6590, 154.7510],\n",
+       "          [152.8370, 424.8210, 150.9870],\n",
+       "          [150.9700, 422.4160, 148.7320],\n",
+       "          [153.5900, 422.6860, 145.9670],\n",
+       "          [157.3050, 423.3410, 145.6600],\n",
+       "          [158.4860, 426.6730, 144.2700],\n",
+       "          [160.7730, 427.0880, 141.2590],\n",
+       "          [163.9630, 429.0530, 141.9600],\n",
+       "          [166.4230, 427.4330, 139.5550],\n",
+       "          [167.6170, 430.8330, 138.3250],\n",
+       "          [170.5180, 431.9210, 140.5320],\n",
+       "          [169.7800, 435.6680, 140.3160],\n",
+       "          [167.6790, 437.8790, 142.6110],\n",
+       "          [167.3350, 435.1910, 145.2890],\n",
+       "          [168.9090, 437.3000, 148.0580],\n",
+       "          [166.1860, 439.9440, 148.5750],\n",
+       "          [164.1010, 438.1300, 151.1890],\n",
+       "          [163.0190, 438.8690, 154.7450],\n",
+       "          [164.6750, 435.6560, 155.9860],\n",
+       "          [168.1170, 436.6550, 154.6330],\n",
+       "          [169.8040, 437.0930, 158.0010],\n",
+       "          [173.2680, 438.4850, 158.5750],\n",
+       "          [174.9270, 440.8490, 156.1460],\n",
+       "          [175.9170, 440.8010, 152.5000],\n",
+       "          [179.6060, 441.0910, 153.4330],\n",
+       "          [180.8000, 439.9830, 156.8710],\n",
+       "          [184.2540, 439.9980, 158.4580],\n",
+       "          [185.5840, 437.2970, 160.7670],\n",
+       "          [186.4140, 439.9000, 163.4480],\n",
+       "          [182.9430, 441.4570, 163.7500],\n",
+       "          [181.4930, 442.0910, 167.2100],\n",
+       "          [178.0780, 440.5360, 167.9060],\n",
+       "          [176.1650, 440.9290, 171.1790],\n",
+       "          [173.0530, 438.9530, 172.1310],\n",
+       "          [170.3990, 440.6990, 174.2170],\n",
+       "          [166.7880, 440.0710, 175.2260],\n",
+       "          [164.4480, 442.1010, 173.0340],\n",
+       "          [160.9780, 443.6110, 173.3810],\n",
+       "          [158.8330, 443.9000, 170.2730],\n",
+       "          [160.1820, 442.9360, 166.8360],\n",
+       "          [163.5010, 444.7230, 166.1630],\n",
+       "          [164.6310, 442.4290, 163.3160],\n",
+       "          [166.6330, 444.2560, 160.6240],\n",
+       "          [166.1700, 447.5210, 162.5220],\n",
+       "          [168.1350, 449.9450, 164.6440],\n",
+       "          [168.0040, 449.0470, 168.3240],\n",
+       "          [168.4210, 451.0540, 171.5220],\n",
+       "          [169.0440, 450.0360, 175.1300],\n",
+       "          [166.1650, 450.3680, 177.6100],\n",
+       "          [166.7800, 449.8300, 181.3340],\n",
+       "          [163.8480, 448.9150, 183.5760],\n",
+       "          [163.3500, 443.9770, 187.9750],\n",
+       "          [165.0240, 442.4930, 184.8870],\n",
+       "          [168.1120, 443.0980, 182.7810],\n",
+       "          [168.1280, 445.9210, 180.2050],\n",
+       "          [166.7230, 445.0130, 176.8030],\n",
+       "          [167.0690, 446.1330, 173.1930],\n",
+       "          [164.0790, 447.8390, 171.5780],\n",
+       "          [163.4630, 449.0630, 168.0350],\n",
+       "          [164.4160, 452.7320, 168.0130],\n",
+       "          [166.7460, 455.3190, 166.5790],\n",
+       "          [167.5460, 459.0090, 166.8440],\n",
+       "          [170.2010, 459.9620, 169.3810],\n",
+       "          [169.5150, 456.8220, 171.4570],\n",
+       "          [170.7440, 454.3880, 168.7860],\n",
+       "          [172.8900, 451.6810, 170.3850],\n",
+       "          [173.2420, 449.3000, 167.4490],\n",
+       "          [171.4700, 447.3670, 164.7150],\n",
+       "          [169.7840, 443.9650, 164.9670],\n",
+       "          [170.8880, 441.4120, 162.3680],\n",
+       "          [169.3320, 438.1640, 163.6510],\n",
+       "          [166.7170, 436.9060, 166.1030],\n",
+       "          [166.2830, 433.5290, 167.7900],\n",
+       "          [162.5490, 433.4260, 168.5040],\n",
+       "          [162.6440, 430.3540, 170.7620],\n",
+       "          [164.9490, 432.1120, 173.2350],\n",
+       "          [163.6530, 435.5760, 172.2210],\n",
+       "          [167.2540, 436.7450, 171.8300],\n",
+       "          [168.3710, 439.3650, 169.3060],\n",
+       "          [171.8910, 439.6680, 167.8930],\n",
+       "          [173.1360, 443.2500, 167.5870],\n",
+       "          [176.0990, 444.7950, 165.8030],\n",
+       "          [177.1670, 447.8660, 167.8270],\n",
+       "          [177.4160, 451.2430, 166.1360],\n",
+       "          [181.1790, 451.0810, 166.7480],\n",
+       "          [181.3560, 448.4350, 163.9980],\n",
+       "          [180.6020, 448.8920, 160.3060],\n",
+       "          [177.8440, 446.8990, 158.6170],\n",
+       "          [176.7540, 446.3160, 155.0150],\n",
+       "          [173.3820, 444.6900, 154.4770],\n",
+       "          [169.6430, 444.9660, 153.9760],\n",
+       "          [167.6080, 446.9450, 156.5140],\n",
+       "          [163.9630, 447.9270, 156.8800],\n",
+       "          [163.1520, 451.2160, 155.1430],\n",
+       "          [160.3870, 453.6570, 156.0420],\n",
+       "          [158.9330, 453.6920, 152.5150],\n",
+       "          [159.8050, 453.2460, 148.8360],\n",
+       "          [161.2560, 456.7540, 148.4440],\n",
+       "          [164.9420, 455.8330, 148.8520],\n",
+       "          [167.3780, 455.9500, 145.9260],\n",
+       "          [171.0380, 455.0610, 145.5060],\n",
+       "          [173.5830, 457.5340, 146.9270],\n",
+       "          [171.2360, 458.9430, 149.5740],\n",
+       "          [171.8150, 459.8640, 153.2130],\n",
+       "          [169.5910, 457.9090, 155.6000],\n",
+       "          [168.4960, 458.3960, 159.2100],\n",
+       "          [167.0870, 455.9020, 161.7110],\n",
+       "          [163.6430, 457.0390, 162.8610],\n",
+       "          [161.8700, 456.3550, 166.1660],\n",
+       "          [160.7490, 452.9260, 164.8950],\n",
+       "          [164.2440, 451.8750, 163.8110],\n",
+       "          [163.4490, 452.3040, 160.1090],\n",
+       "          [165.6820, 454.0140, 157.5590],\n",
+       "          [164.2940, 457.2310, 156.0670],\n",
+       "          [165.6810, 459.7290, 153.5780],\n",
+       "          [167.3870, 462.7790, 155.0760],\n",
+       "          [170.2060, 465.8340, 161.8360],\n",
+       "          [172.1620, 462.7680, 160.6890],\n",
+       "          [174.0280, 461.7040, 163.8180],\n",
+       "          [175.2110, 458.5660, 161.9930],\n",
+       "          [176.5410, 458.2410, 158.4450],\n",
+       "          [174.5550, 455.7530, 156.3540],\n",
+       "          [174.3580, 455.6030, 152.5600],\n",
+       "          [171.8590, 453.6260, 150.4910],\n",
+       "          [173.6300, 451.2730, 148.0810],\n",
+       "          [170.4820, 450.1750, 146.2000],\n",
+       "          [166.8520, 451.1980, 145.7810],\n",
+       "          [164.0730, 450.1880, 148.1490],\n",
+       "          [162.6710, 446.7750, 147.1980],\n",
+       "          [159.1420, 445.8080, 148.2240],\n",
+       "          [158.3160, 442.3640, 149.5950],\n",
+       "          [155.0300, 443.1470, 151.4050],\n",
+       "          [152.6800, 446.0960, 151.8500],\n",
+       "          [155.0610, 447.4890, 154.4910],\n",
+       "          [158.2010, 445.3460, 154.0190],\n",
+       "          [160.8090, 447.5340, 152.2870],\n",
+       "          [164.4130, 446.3150, 152.1360],\n",
+       "          [167.3360, 448.5380, 151.1380],\n",
+       "          [171.0310, 447.6630, 150.9500],\n",
+       "          [172.8160, 450.1440, 153.2260],\n",
+       "          [176.3550, 450.4540, 154.5850],\n",
+       "          [176.6310, 452.0780, 158.0240],\n",
+       "          [180.1600, 453.0970, 159.0310],\n",
+       "          [179.9110, 454.9710, 162.3160],\n",
+       "          [178.7370, 458.4030, 163.4610],\n",
+       "          [179.1350, 461.9150, 162.0830],\n",
+       "          [180.3160, 464.9940, 163.9950],\n",
+       "          [195.1700, 454.0210, 194.1160],\n",
+       "          [196.4220, 450.5570, 193.1610],\n",
+       "          [195.0530, 447.5260, 191.3340],\n",
+       "          [195.7530, 445.4320, 194.4410],\n",
+       "          [195.6760, 446.8960, 197.9710],\n",
+       "          [199.1230, 447.9300, 199.1640],\n",
+       "          [198.4070, 447.0140, 202.7840],\n",
+       "          [198.1130, 443.3080, 203.5090],\n",
+       "          [195.0350, 441.7530, 205.1060],\n",
+       "          [194.7040, 440.1720, 208.5300],\n",
+       "          [191.7420, 437.9460, 207.6670],\n",
+       "          [190.9460, 435.0990, 205.3140],\n",
+       "          [188.0870, 435.6760, 202.8880],\n",
+       "          [184.8300, 433.7320, 203.1830],\n",
+       "          [184.1270, 431.8260, 199.9570],\n",
+       "          [180.5450, 430.8500, 200.7460],\n",
+       "          [178.2470, 433.3320, 199.0220],\n",
+       "          [178.2770, 433.9260, 195.2680],\n",
+       "          [176.3170, 437.2030, 195.0280],\n",
+       "          [177.7690, 440.7330, 194.9660],\n",
+       "          [181.3350, 439.3910, 195.1330],\n",
+       "          [182.5820, 441.3460, 192.0940],\n",
+       "          [182.3530, 444.9560, 193.3610],\n",
+       "          [186.0870, 445.2240, 194.0290],\n",
+       "          [188.7230, 447.5570, 192.6190],\n",
+       "          [190.8760, 444.5870, 191.5490],\n",
+       "          [188.2670, 443.5230, 188.9590],\n",
+       "          [190.1360, 444.3010, 185.7360],\n",
+       "          [188.8140, 444.1990, 182.1980],\n",
+       "          [185.1300, 444.1560, 181.3470],\n",
+       "          [182.1290, 442.2030, 182.6060],\n",
+       "          [181.3330, 441.2700, 178.9860],\n",
+       "          [184.3690, 440.6570, 176.7700],\n",
+       "          [184.4240, 439.7520, 173.0740],\n",
+       "          [187.1540, 437.5560, 171.6020],\n",
+       "          [187.5170, 439.9750, 168.6500],\n",
+       "          [188.4180, 443.0500, 170.7240],\n",
+       "          [191.4360, 445.2100, 169.8910],\n",
+       "          [193.7480, 445.6510, 172.8940],\n",
+       "          [196.8200, 447.8990, 173.0280],\n",
+       "          [199.4910, 448.6120, 175.6310],\n",
+       "          [201.2080, 451.9090, 176.3780],\n",
+       "          [203.6110, 453.4740, 178.8680],\n",
+       "          [201.6130, 454.9210, 181.7570],\n",
+       "          [202.0690, 458.3840, 183.2730],\n",
+       "          [200.4970, 458.3690, 186.7180],\n",
+       "          [197.7960, 455.8130, 187.5240],\n",
+       "          [195.1460, 455.8470, 184.7720],\n",
+       "          [193.6160, 452.5740, 185.9790],\n",
+       "          [189.8600, 452.3800, 185.3150],\n",
+       "          [189.9440, 455.8580, 183.7680],\n",
+       "          [189.7400, 457.4310, 180.3410],\n",
+       "          [193.1630, 457.9480, 178.8100],\n",
+       "          [194.6770, 460.5170, 176.4640],\n",
+       "          [197.9720, 460.6430, 174.6010],\n",
+       "          [200.7710, 462.5710, 176.3210],\n",
+       "          [203.7970, 463.7400, 174.3330],\n",
+       "          [207.1150, 464.5820, 175.9770],\n",
+       "          [210.1390, 466.5260, 174.8060],\n",
+       "          [210.8850, 463.9390, 172.1200],\n",
+       "          [213.6440, 461.8050, 173.6150],\n",
+       "          [211.1280, 459.2430, 174.9140],\n",
+       "          [208.0390, 457.7610, 173.2930],\n",
+       "          [204.4770, 458.9110, 173.9310],\n",
+       "          [202.5240, 457.8000, 177.0030],\n",
+       "          [198.9190, 457.7220, 178.2440],\n",
+       "          [197.5730, 459.9040, 181.0500],\n",
+       "          [194.2310, 460.2350, 182.8060],\n",
+       "          [192.1140, 462.6430, 180.7850],\n",
+       "          [189.0250, 463.1720, 178.7000],\n",
+       "          [187.1240, 465.8360, 176.8220],\n",
+       "          [188.2850, 466.3320, 173.2420],\n",
+       "          [191.5580, 464.4520, 173.8600],\n",
+       "          [189.9950, 461.1280, 174.9010],\n",
+       "          [191.8440, 458.2970, 173.1540],\n",
+       "          [190.4670, 455.2690, 174.9900],\n",
+       "          [189.7380, 453.5820, 178.3070],\n",
+       "          [192.3390, 451.8830, 180.5000],\n",
+       "          [191.2020, 448.5340, 181.9050],\n",
+       "          [194.3800, 447.0280, 183.4050],\n",
+       "          [197.8700, 447.9650, 184.5850],\n",
+       "          [200.9840, 445.8270, 184.9660],\n",
+       "          [202.8290, 447.6220, 187.7590],\n",
+       "          [206.2120, 445.9180, 187.3400],\n",
+       "          [206.6040, 447.1310, 183.7460],\n",
+       "          [204.4430, 450.2380, 184.1830],\n",
+       "          [202.2310, 449.1370, 181.2890],\n",
+       "          [198.6660, 450.3790, 180.7720],\n",
+       "          [196.2310, 448.2490, 178.7680],\n",
+       "          [194.0550, 450.4240, 176.5260],\n",
+       "          [190.8490, 449.7090, 174.6070],\n",
+       "          [190.4360, 452.3080, 171.8240],\n",
+       "          [187.3830, 454.5590, 171.7980],\n",
+       "          [186.4670, 452.9730, 168.4480],\n",
+       "          [185.9050, 449.6480, 170.2580],\n",
+       "          [182.8890, 448.8910, 172.4330],\n",
+       "          [183.5680, 448.6030, 176.1630],\n",
+       "          [181.3330, 447.5400, 179.0590],\n",
+       "          [182.7340, 447.5270, 182.5670],\n",
+       "          [183.3280, 449.2260, 185.8920],\n",
+       "          [184.8760, 452.6980, 185.6990],\n",
+       "          [185.7140, 455.4170, 188.2030],\n",
+       "          [182.9420, 457.9530, 188.8300],\n",
+       "          [183.2930, 461.5650, 189.9350],\n",
+       "          [182.7580, 462.1320, 193.6660],\n",
+       "          [180.0620, 464.7490, 193.1030],\n",
+       "          [176.4120, 463.9030, 192.5790],\n",
+       "          [177.1550, 460.1730, 192.6930],\n",
+       "          [173.8180, 459.5980, 194.4390],\n",
+       "          [172.0770, 461.1030, 191.4040],\n",
+       "          [173.6950, 458.5410, 189.0910],\n",
+       "          [171.0060, 456.0190, 188.1270],\n",
+       "          [170.8210, 453.3390, 185.4210],\n",
+       "          [170.0180, 454.8270, 182.0350],\n",
+       "          [171.8690, 458.1080, 182.6590],\n",
+       "          [174.2320, 459.6520, 180.1210],\n",
+       "          [177.8640, 459.9130, 181.2230],\n",
+       "          [180.8780, 461.8410, 179.9490],\n",
+       "          [184.5650, 461.8070, 180.8430],\n",
+       "          [186.0650, 464.7990, 182.6440],\n",
+       "          [189.5780, 466.2740, 182.6170],\n",
+       "          [190.7120, 463.6340, 185.1390],\n",
+       "          [189.2980, 460.6630, 183.2290],\n",
+       "          [186.3560, 460.3000, 185.6240],\n",
+       "          [182.7840, 459.6590, 184.5210],\n",
+       "          [180.2130, 462.3320, 185.3590],\n",
+       "          [176.5320, 462.8410, 184.6140],\n",
+       "          [175.8580, 465.0300, 181.5790],\n",
+       "          [178.1930, 465.7110, 174.9000],\n",
+       "          [180.9150, 464.1190, 172.7670],\n",
+       "          [181.2510, 460.5420, 174.1130],\n",
+       "          [178.4430, 457.9660, 174.2140],\n",
+       "          [178.5420, 456.4790, 177.7180],\n",
+       "          [175.4730, 455.1590, 179.5500],\n",
+       "          [175.4670, 454.0690, 183.1870],\n",
+       "          [174.1210, 450.5430, 183.6830],\n",
+       "          [173.9850, 450.6340, 187.5030],\n",
+       "          [174.3640, 452.9240, 190.5240],\n",
+       "          [177.5140, 454.2200, 192.1820],\n",
+       "          [179.3060, 451.6570, 194.3540],\n",
+       "          [181.4890, 453.2110, 197.0570],\n",
+       "          [184.7390, 451.3110, 197.6330],\n",
+       "          [186.4870, 454.0230, 199.6690],\n",
+       "          [185.5270, 457.3810, 201.1790],\n",
+       "          [186.7910, 459.0050, 197.9610],\n",
+       "          [186.6010, 456.0430, 195.5320],\n",
+       "          [183.3680, 455.3620, 193.6270],\n",
+       "          [182.9110, 452.9430, 190.7280],\n",
+       "          [180.0190, 452.5230, 188.3060],\n",
+       "          [179.0910, 450.0250, 185.5960],\n",
+       "          [178.9850, 451.7840, 182.2280],\n",
+       "          [179.0510, 451.0080, 178.5100],\n",
+       "          [181.0530, 453.1550, 176.0780],\n",
+       "          [180.1320, 452.7170, 172.4060],\n",
+       "          [182.0570, 455.4530, 170.6120],\n",
+       "          [181.8140, 459.1820, 170.0770],\n",
+       "          [178.5990, 461.0580, 169.3130],\n",
+       "          [177.7620, 463.8650, 166.9020]]),\n",
        "  'residue_ids': tensor([  4.,   4.,   4.,  ..., 182., 182., 182.]),\n",
        "  'chain_ids': tensor([0., 0., 0.,  ..., 1., 1., 1.])},\n",
        " 'id': '8phr__X4_UNDEFINED--8phr__W4_UNDEFINED',\n",
@@ -906,7 +1568,7 @@
        " 'target_id': '8phr__X4_UNDEFINED-R--8phr__W4_UNDEFINED-L'}"
       ]
      },
-     "execution_count": 3,
+     "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -918,7 +1580,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 16,
    "id": "4a6f64b4-f027-4286-b444-68cd9f6b59a4",
    "metadata": {},
    "outputs": [
@@ -928,7 +1590,7 @@
        "torch.Size([1316, 3])"
       ]
      },
-     "execution_count": 10,
+     "execution_count": 16,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -945,7 +1607,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 17,
    "id": "01025b2b-1962-43b3-8c4e-a33b6a3d6231",
    "metadata": {},
    "outputs": [
@@ -956,7 +1618,7 @@
       "Help on class PinderDataset in module pinder.core.loader.dataset:\n",
       "\n",
       "class PinderDataset(torch.utils.data.dataset.Dataset)\n",
-      " |  PinderDataset(split: 'str | None' = None, index: 'pd.DataFrame | None' = None, metadata: 'pd.DataFrame | None' = None, monomer_priority: 'str' = 'holo', base_filters: 'list[PinderFilterBase]' = [], sub_filters: 'list[PinderFilterSubBase]' = [], structure_filters: 'list[StructureFilter]' = [], structure_transforms_target: 'list[StructureTransform]' = [], structure_transforms_feature: 'list[StructureTransform]' = [], transform: 'Callable[[Structure], torch.Tensor | dict[str, torch.Tensor]]' = <function structure2tensor_transform at 0x1e6b9a4d0>, target_transform: 'Callable[[Structure], torch.Tensor | dict[str, torch.Tensor]]' = <function structure2tensor_transform at 0x1e6b9a4d0>, ids: 'list[str] | None' = None, fallback_to_holo: 'bool' = True, use_canonical_apo: 'bool' = True, crop_equal_monomer_shapes: 'bool' = True, index_query: 'str | None' = None, metadata_query: 'str | None' = None, pre_specified_monomers: 'dict[str, str] | pd.DataFrame | None' = None, **kwargs: 'Any') -> 'None'\n",
+      " |  PinderDataset(split: 'str | None' = None, index: 'pd.DataFrame | None' = None, metadata: 'pd.DataFrame | None' = None, monomer_priority: 'str' = 'holo', base_filters: 'list[PinderFilterBase]' = [], sub_filters: 'list[PinderFilterSubBase]' = [], structure_filters: 'list[StructureFilter]' = [], structure_transforms_target: 'list[StructureTransform]' = [], structure_transforms_feature: 'list[StructureTransform]' = [], transform: 'Callable[[Structure], torch.Tensor | dict[str, torch.Tensor]]' = <function structure2tensor_transform at 0x1f9f84d30>, target_transform: 'Callable[[Structure], torch.Tensor | dict[str, torch.Tensor]]' = <function structure2tensor_transform at 0x1f9f84d30>, ids: 'list[str] | None' = None, fallback_to_holo: 'bool' = True, use_canonical_apo: 'bool' = True, crop_equal_monomer_shapes: 'bool' = True, index_query: 'str | None' = None, metadata_query: 'str | None' = None, pre_specified_monomers: 'dict[str, str] | pd.DataFrame | None' = None, **kwargs: 'Any') -> 'None'\n",
       " |  \n",
       " |  Method resolution order:\n",
       " |      PinderDataset\n",
@@ -968,7 +1630,7 @@
       " |  \n",
       " |  __getitem__(self, idx: 'int') -> 'dict[str, dict[str, torch.Tensor] | torch.Tensor]'\n",
       " |  \n",
-      " |  __init__(self, split: 'str | None' = None, index: 'pd.DataFrame | None' = None, metadata: 'pd.DataFrame | None' = None, monomer_priority: 'str' = 'holo', base_filters: 'list[PinderFilterBase]' = [], sub_filters: 'list[PinderFilterSubBase]' = [], structure_filters: 'list[StructureFilter]' = [], structure_transforms_target: 'list[StructureTransform]' = [], structure_transforms_feature: 'list[StructureTransform]' = [], transform: 'Callable[[Structure], torch.Tensor | dict[str, torch.Tensor]]' = <function structure2tensor_transform at 0x1e6b9a4d0>, target_transform: 'Callable[[Structure], torch.Tensor | dict[str, torch.Tensor]]' = <function structure2tensor_transform at 0x1e6b9a4d0>, ids: 'list[str] | None' = None, fallback_to_holo: 'bool' = True, use_canonical_apo: 'bool' = True, crop_equal_monomer_shapes: 'bool' = True, index_query: 'str | None' = None, metadata_query: 'str | None' = None, pre_specified_monomers: 'dict[str, str] | pd.DataFrame | None' = None, **kwargs: 'Any') -> 'None'\n",
+      " |  __init__(self, split: 'str | None' = None, index: 'pd.DataFrame | None' = None, metadata: 'pd.DataFrame | None' = None, monomer_priority: 'str' = 'holo', base_filters: 'list[PinderFilterBase]' = [], sub_filters: 'list[PinderFilterSubBase]' = [], structure_filters: 'list[StructureFilter]' = [], structure_transforms_target: 'list[StructureTransform]' = [], structure_transforms_feature: 'list[StructureTransform]' = [], transform: 'Callable[[Structure], torch.Tensor | dict[str, torch.Tensor]]' = <function structure2tensor_transform at 0x1f9f84d30>, target_transform: 'Callable[[Structure], torch.Tensor | dict[str, torch.Tensor]]' = <function structure2tensor_transform at 0x1f9f84d30>, ids: 'list[str] | None' = None, fallback_to_holo: 'bool' = True, use_canonical_apo: 'bool' = True, crop_equal_monomer_shapes: 'bool' = True, index_query: 'str | None' = None, metadata_query: 'str | None' = None, pre_specified_monomers: 'dict[str, str] | pd.DataFrame | None' = None, **kwargs: 'Any') -> 'None'\n",
       " |      Initialize self.  See help(type(self)) for accurate signature.\n",
       " |  \n",
       " |  __len__(self) -> 'int'\n",
@@ -1039,7 +1701,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 18,
    "id": "fcb50bed-7e0a-4d4b-a1d7-6bc003161aab",
    "metadata": {},
    "outputs": [
@@ -1047,10 +1709,11 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2024-09-23 15:44:49,788 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=8, items=8\n",
-      "2024-09-23 15:44:50,077 | pinder.core.utils.cloud.process_many:23 | INFO : runtime succeeded:      0.29s\n",
-      "2024-09-23 15:44:50,350 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=5, items=5\n",
-      "2024-09-23 15:44:50,645 | pinder.core.utils.cloud.process_many:23 | INFO : runtime succeeded:      0.30s\n"
+      "2024-09-25 16:23:26,740 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=5, items=5\n",
+      "2024-09-25 16:23:27,432 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=5, items=5\n",
+      "2024-09-25 16:23:27,691 | pinder.core.loader.filters:436 | WARNING : Skipping FilterSubRmsds on 3muw__B33_P03316--3muw__B6_P03316 as holo monomers have been filtered out\n",
+      "2024-09-25 16:23:27,789 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=5, items=5\n",
+      "2024-09-25 16:23:28,201 | pinder.core.loader.structure:595 | ERROR : no common residues found! 2in2__A1_P03303--2rs3__C39_P03303-L\n"
      ]
     }
    ],
@@ -1110,7 +1773,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 27,
    "id": "0de517fc-6e7f-4682-af48-980d27638b21",
    "metadata": {},
    "outputs": [
@@ -1119,7 +1782,7 @@
      "output_type": "stream",
      "text": [
       "Processing...\n",
-      "2024-09-23 15:45:00,968 | pinder.core.loader.dataset:541 | INFO : Finished processing, only 5 systems\n",
+      "2024-09-25 16:29:35,300 | pinder.core.loader.dataset:543 | INFO : Finished processing, only 5 systems\n",
       "Done!\n"
      ]
     },
@@ -1129,7 +1792,7 @@
        "PPIDataset(5)"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 27,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1157,7 +1820,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 28,
    "id": "5392a214-7426-4be2-bc81-a00702c82c55",
    "metadata": {},
    "outputs": [
@@ -1178,12 +1841,12 @@
        "    chain=[1],\n",
        "  },\n",
        "  ligand_atom={\n",
-       "    x=[1198, 12],\n",
+       "    x=[1198, 1],\n",
        "    pos=[1198, 3],\n",
        "    edge_index=[2, 11980],\n",
        "  },\n",
        "  receptor_atom={\n",
-       "    x=[1358, 12],\n",
+       "    x=[1358, 1],\n",
        "    pos=[1358, 3],\n",
        "    edge_index=[2, 13580],\n",
        "  },\n",
@@ -1194,7 +1857,7 @@
        ")"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 28,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1220,7 +1883,66 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 29,
+   "id": "7948dcbd-fcb9-422e-89be-fbf3d1c4f13c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'ligand_residue': 0,\n",
+       " 'receptor_residue': 0,\n",
+       " 'ligand_atom': 1,\n",
+       " 'receptor_atom': 1,\n",
+       " 'pdb': 0}"
+      ]
+     },
+     "execution_count": 29,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "data_item.num_node_features"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "id": "ff827dac-c7e5-4e81-99ff-ff9e67c3cbbe",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'x': tensor([[0.],\n",
+       "        [1.],\n",
+       "        [2.],\n",
+       "        ...,\n",
+       "        [2.],\n",
+       "        [3.],\n",
+       "        [8.]]), 'pos': tensor([[194.1170, 453.4250, 194.9300],\n",
+       "        [195.1700, 454.0210, 194.1160],\n",
+       "        [196.0630, 452.9460, 193.5050],\n",
+       "        ...,\n",
+       "        [177.4130, 465.0800, 167.7550],\n",
+       "        [176.8000, 464.9490, 168.8150],\n",
+       "        [176.7910, 463.7350, 165.7380]]), 'edge_index': tensor([[   1,    4,    2,  ..., 1186, 1187, 1182],\n",
+       "        [   0,    0,    0,  ..., 1197, 1197, 1197]])}"
+      ]
+     },
+     "execution_count": 30,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "data_item[\"ligand_atom\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
    "id": "6fa927a9-cde9-4f32-acf9-f5c989b32966",
    "metadata": {},
    "outputs": [
@@ -1241,12 +1963,12 @@
        "    chain=[1],\n",
        "  },\n",
        "  ligand_atom={\n",
-       "    x=[1198, 12],\n",
+       "    x=[1198, 1],\n",
        "    pos=[1198, 3],\n",
        "    edge_index=[2, 11980],\n",
        "  },\n",
        "  receptor_atom={\n",
-       "    x=[1358, 12],\n",
+       "    x=[1358, 1],\n",
        "    pos=[1358, 3],\n",
        "    edge_index=[2, 13580],\n",
        "  },\n",
@@ -1257,7 +1979,7 @@
        ")"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 31,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1271,7 +1993,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 32,
    "id": "75e13a76-cf9f-47dd-90c2-370c0652074f",
    "metadata": {},
    "outputs": [
@@ -1327,7 +2049,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 33,
    "id": "ce872929-1706-4a09-ad67-63457511099f",
    "metadata": {
     "scrolled": true
@@ -1350,19 +2072,19 @@
        "    chain=[1],\n",
        "  },\n",
        "  ligand_atom={\n",
-       "    x=[1032, 12],\n",
+       "    x=[1032, 1],\n",
        "    pos=[1032, 3],\n",
        "    edge_index=[2, 10320],\n",
        "  },\n",
        "  receptor_atom={\n",
-       "    x=[1441, 12],\n",
+       "    x=[1441, 1],\n",
        "    pos=[1441, 3],\n",
        "    edge_index=[2, 14410],\n",
        "  }\n",
        ")"
       ]
      },
-     "execution_count": 10,
+     "execution_count": 33,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1389,8 +2111,8 @@
     "expected_num_feats = {\n",
     "    'ligand_residue': 0,\n",
     "    'receptor_residue': 0,\n",
-    "    'ligand_atom': 12,\n",
-    "    'receptor_atom': 12\n",
+    "    'ligand_atom': 1,\n",
+    "    'receptor_atom': 1\n",
     "}\n",
     "for k, v in expected_num_feats.items():\n",
     "    assert holo_data.num_node_features[k] == v\n",
@@ -1404,7 +2126,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 34,
    "id": "b2913a8f-2992-46de-ab74-1fca594d6088",
    "metadata": {},
    "outputs": [
@@ -1425,19 +2147,19 @@
        "    chain=[1],\n",
        "  },\n",
        "  ligand_atom={\n",
-       "    x=[1350, 12],\n",
+       "    x=[1350, 1],\n",
        "    pos=[1350, 3],\n",
        "    edge_index=[2, 13500],\n",
        "  },\n",
        "  receptor_atom={\n",
-       "    x=[1710, 12],\n",
+       "    x=[1710, 1],\n",
        "    pos=[1710, 3],\n",
        "    edge_index=[2, 17100],\n",
        "  }\n",
        ")"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 34,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1457,8 +2179,8 @@
     "expected_num_feats = {\n",
     "    'ligand_residue': 0,\n",
     "    'receptor_residue': 0,\n",
-    "    'ligand_atom': 12,\n",
-    "    'receptor_atom': 12\n",
+    "    'ligand_atom': 1,\n",
+    "    'receptor_atom': 1,\n",
     "}\n",
     "for k, v in expected_num_feats.items():\n",
     "    assert apo_data.num_node_features[k] == v\n",
@@ -1483,17 +2205,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 35,
    "id": "29ca7612-8a1b-4cf7-a790-cc2b5b26e7e5",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "<torch_geometric.loader.dataloader.DataLoader at 0x1fb60f550>"
+       "<torch_geometric.loader.dataloader.DataLoader at 0x20265ebf0>"
       ]
      },
-     "execution_count": 12,
+     "execution_count": 35,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1516,7 +2238,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 36,
    "id": "66f9a0ef-1c1f-47bd-ba86-f2c3340c34ec",
    "metadata": {},
    "outputs": [
@@ -1526,7 +2248,7 @@
        "PPIDataset(5)"
       ]
      },
-     "execution_count": 13,
+     "execution_count": 36,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1572,7 +2294,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 37,
    "id": "bb468005-5ffb-45e8-9f03-f073cb6f7527",
    "metadata": {},
    "outputs": [],
@@ -1696,7 +2418,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 38,
    "id": "def48617-4288-452f-af9e-55f3289d7788",
    "metadata": {},
    "outputs": [
@@ -1719,7 +2441,7 @@
        " ))"
       ]
      },
-     "execution_count": 15,
+     "execution_count": 38,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1744,7 +2466,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 39,
    "id": "c13a072b-ca2b-47bb-9ed1-2a88885aad93",
    "metadata": {},
    "outputs": [],
@@ -1755,20 +2477,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 40,
    "id": "84f1a37c-4ac5-4086-bfc8-260bc8f2de8d",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "(array([[-19.331,  10.568,   4.591],\n",
-       "        [-18.082,  10.305,   3.883],\n",
-       "        [-16.867,  10.692,   4.727],\n",
+       "(array([[-12.6210985,  -9.128864 ,  17.258345 ],\n",
+       "        [-13.660538 ,  -8.840154 ,  16.265753 ],\n",
+       "        [-13.38884  ,  -7.5286074,  15.524057 ],\n",
        "        ...,\n",
-       "        [ 13.363,  11.96 ,  -6.092],\n",
-       "        [ 15.081,   9.028,  -5.759],\n",
-       "        [ 12.459,  10.17 ,  -6.767]], dtype=float32),\n",
+       "        [  7.1754823, -19.776093 ,  21.191608 ],\n",
+       "        [  9.592421 , -19.601936 ,  19.197937 ],\n",
+       "        [  9.432583 , -17.757698 ,  20.458267 ]], dtype=float32),\n",
        " array([[-13.215324 , -11.076905 ,  15.214827 ],\n",
        "        [-14.133494 , -10.301853 ,  14.386162 ],\n",
        "        [-13.614427 ,  -8.882867 ,  14.150835 ],\n",
@@ -1778,7 +2500,7 @@
        "        [  7.0213795, -15.329907 ,  18.152586 ]], dtype=float32))"
       ]
      },
-     "execution_count": 17,
+     "execution_count": 40,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1800,17 +2522,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 41,
    "id": "8c495bfe-3aa8-4544-b94e-3bcb6e149dba",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "<torch.utils.data.sampler.WeightedRandomSampler at 0x1fb60c5b0>"
+       "<torch.utils.data.sampler.WeightedRandomSampler at 0x200146500>"
       ]
      },
-     "execution_count": 18,
+     "execution_count": 41,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1855,7 +2577,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 42,
    "id": "cc7f7dc6-f01b-4d52-a4b5-e0dd19d68fe4",
    "metadata": {},
    "outputs": [
@@ -1863,30 +2585,29 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2024-09-23 15:45:18,411 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=2, items=2\n",
-      "2024-09-23 15:45:18,707 | pinder.core.utils.cloud.process_many:23 | INFO : runtime succeeded:      0.30s\n"
+      "2024-09-25 16:31:53,018 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=1, items=1\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "(tensor([[[37.2095,  8.0406, 72.5705],\n",
-       "          [36.1588,  7.1319, 72.1171],\n",
-       "          [35.5337,  7.6365, 70.8136],\n",
+       "(tensor([[[-44.9000, -13.0073,   3.9172],\n",
+       "          [-44.8239, -14.3579,   3.3580],\n",
+       "          [-44.6508, -14.3793,   1.8236],\n",
        "          ...,\n",
-       "          [40.5402,  3.1015, 96.2705],\n",
-       "          [39.7635,  4.0274, 94.4056],\n",
-       "          [39.5253,  3.4914, 95.5608]]]),\n",
-       " tensor([[[29.6114,  7.0679, 73.2040],\n",
-       "          [30.3889,  6.7263, 72.0241],\n",
-       "          [31.5403,  7.6852, 71.7204],\n",
+       "          [-41.4551, -36.3519,   2.2509],\n",
+       "          [-42.4746, -36.6704,   1.1575],\n",
+       "          [-43.0579, -35.4178,   0.6238]]]),\n",
+       " tensor([[[-47.9769,  -9.1210,   3.2415],\n",
+       "          [-47.0982,  -9.0194,   2.0802],\n",
+       "          [-46.2697, -10.2855,   1.8961],\n",
        "          ...,\n",
-       "          [35.5366, -4.1661, 95.5403],\n",
-       "          [34.4548, -4.3722, 96.2721],\n",
-       "          [35.4055, -4.2196, 94.2242]]]))"
+       "          [-35.2793, -42.3302,   2.4468],\n",
+       "          [-34.2672, -43.3448,   1.9357],\n",
+       "          [-34.0432, -43.2068,   0.4709]]]))"
       ]
      },
-     "execution_count": 19,
+     "execution_count": 42,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1917,17 +2638,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 43,
    "id": "fe172bdd-09eb-45ed-bc51-28202e5158b6",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "<torch.utils.data.dataloader.DataLoader at 0x1fb6404c0>"
+       "<torch.utils.data.dataloader.DataLoader at 0x1ff3c7850>"
       ]
      },
-     "execution_count": 20,
+     "execution_count": 43,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1974,7 +2695,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 44,
    "id": "b7602385-2893-4c09-afc3-6a10433c7ada",
    "metadata": {},
    "outputs": [
@@ -1982,30 +2703,29 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2024-09-23 15:45:25,616 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=2, items=2\n",
-      "2024-09-23 15:45:25,901 | pinder.core.utils.cloud.process_many:23 | INFO : runtime succeeded:      0.28s\n"
+      "2024-09-25 16:31:59,644 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=10, items=15\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "(tensor([[[226.8320, 169.1410, 206.0350],\n",
-       "          [226.2280, 168.7600, 207.3040],\n",
-       "          [227.3330, 168.7310, 208.3410],\n",
+       "(tensor([[[ 70.4514,  -6.9500,  29.3370],\n",
+       "          [ 69.9283,  -6.6501,  30.6817],\n",
+       "          [ 68.9000,  -5.5045,  30.7376],\n",
        "          ...,\n",
-       "          [205.7270, 143.0570, 181.4730],\n",
-       "          [206.4940, 142.1170, 181.2560],\n",
-       "          [204.8250, 144.8850, 182.9280]]]),\n",
-       " tensor([[[226.8320, 169.1410, 206.0350],\n",
-       "          [226.2280, 168.7600, 207.3040],\n",
-       "          [227.3330, 168.7310, 208.3410],\n",
+       "          [ 67.3974, -52.2529,  39.1974],\n",
+       "          [ 65.9296, -51.8409,  39.2738],\n",
+       "          [ 68.1239, -51.1915,  38.3785]]]),\n",
+       " tensor([[[ 68.1820,  -6.3490,  28.9990],\n",
+       "          [ 68.8730,  -5.0680,  28.9240],\n",
+       "          [ 68.3080,  -4.1750,  27.8270],\n",
        "          ...,\n",
-       "          [205.7270, 143.0570, 181.4730],\n",
-       "          [206.4940, 142.1170, 181.2560],\n",
-       "          [204.8250, 144.8850, 182.9280]]]))"
+       "          [ 68.4140, -52.5670,  40.9670],\n",
+       "          [ 67.2560, -52.0200,  41.7820],\n",
+       "          [ 69.1350, -51.4480,  40.2400]]]))"
       ]
      },
-     "execution_count": 21,
+     "execution_count": 44,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -2035,7 +2755,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": 45,
    "id": "cce6f303-99e7-4d42-805d-72f39d9c61b4",
    "metadata": {},
    "outputs": [
@@ -2074,7 +2794,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 46,
    "id": "4fbf640e-3348-439c-a8d0-ae4bb2c86dc3",
    "metadata": {},
    "outputs": [
@@ -2082,48 +2802,45 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2024-09-23 15:45:27,403 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=10, items=169\n",
-      "2024-09-23 15:45:30,832 | pinder.core.utils.cloud.process_many:23 | INFO : runtime succeeded:      3.43s\n",
-      "2024-09-23 15:45:31,248 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=5, items=5\n",
-      "2024-09-23 15:45:31,445 | pinder.core.utils.cloud.process_many:23 | INFO : runtime succeeded:      0.20s\n"
+      "2024-09-25 16:32:03,705 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=5, items=5\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "(tensor([[[ -20.7939,   16.2259,   28.9029],\n",
-       "          [ -20.3926,   16.6965,   27.5845],\n",
-       "          [ -19.1855,   15.8943,   27.0868],\n",
-       "          ...,\n",
-       "          [  -3.5845,    7.8611,    8.0106],\n",
-       "          [  -6.1668,    6.1222,    8.1268],\n",
-       "          [  -5.3776,    5.6234,    9.3411]],\n",
-       " \n",
-       "         [[   5.1565,   -8.1926,   18.2980],\n",
-       "          [   4.2819,   -7.6345,   19.3660],\n",
-       "          [   3.1236,   -8.5983,   19.6553],\n",
+       "(tensor([[[ 239.4760,  376.4950,  244.9000],\n",
+       "          [ 238.9000,  377.8010,  245.2000],\n",
+       "          [ 237.8190,  377.6950,  246.2690],\n",
        "          ...,\n",
        "          [-100.0000, -100.0000, -100.0000],\n",
        "          [-100.0000, -100.0000, -100.0000],\n",
-       "          [-100.0000, -100.0000, -100.0000]]]),\n",
-       " tensor([[[ -20.9630,   15.4670,   28.5250],\n",
-       "          [ -20.3050,   16.5060,   27.7270],\n",
-       "          [ -18.9740,   15.9710,   27.1630],\n",
-       "          ...,\n",
-       "          [  -0.7240,   -5.2800,    4.0790],\n",
-       "          [  -0.9060,   -6.4480,    4.9780],\n",
-       "          [  -2.0090,   -6.1390,    5.9730]],\n",
+       "          [-100.0000, -100.0000, -100.0000]],\n",
        " \n",
-       "         [[   5.0550,   -8.1610,   18.0120],\n",
-       "          [   4.3920,   -7.8030,   19.2620],\n",
-       "          [   3.3010,   -8.8210,   19.5290],\n",
+       "         [[  31.4372, -103.9317,   46.1102],\n",
+       "          [  32.4872, -104.3322,   45.1591],\n",
+       "          [  32.0485, -104.2013,   43.6919],\n",
+       "          ...,\n",
+       "          [  10.1052, -129.8667,   21.4974],\n",
+       "          [  13.3537, -129.0618,   23.8115],\n",
+       "          [  12.3266, -127.9412,   23.6833]]]),\n",
+       " tensor([[[ 239.4760,  376.4950,  244.9000],\n",
+       "          [ 238.9000,  377.8010,  245.2000],\n",
+       "          [ 237.8190,  377.6950,  246.2690],\n",
        "          ...,\n",
        "          [-100.0000, -100.0000, -100.0000],\n",
        "          [-100.0000, -100.0000, -100.0000],\n",
-       "          [-100.0000, -100.0000, -100.0000]]]))"
+       "          [-100.0000, -100.0000, -100.0000]],\n",
+       " \n",
+       "         [[  31.9020, -105.5020,   46.4300],\n",
+       "          [  31.2640, -104.3600,   45.7100],\n",
+       "          [  31.5600, -104.3840,   44.1990],\n",
+       "          ...,\n",
+       "          [   8.6080, -130.5800,   20.6600],\n",
+       "          [   7.2110, -130.2610,   21.1240],\n",
+       "          [   7.2070, -128.8110,   21.6200]]]))"
       ]
      },
-     "execution_count": 23,
+     "execution_count": 46,
      "metadata": {},
      "output_type": "execute_result"
     }
diff --git a/examples/pinder-mlsb.ipynb b/examples/pinder-mlsb.ipynb
index fde36cd..a51a6dd 100644
--- a/examples/pinder-mlsb.ipynb
+++ b/examples/pinder-mlsb.ipynb
@@ -232,20 +232,20 @@
     {
      "data": {
       "text/plain": [
-       "({'native': 'pdbs/8ir7__S1_P51765--8ir7__AA1_Q7DGD4.pdb',\n",
-       "  'holo_R': 'pdbs/8ir7__S1_P51765-R.pdb',\n",
-       "  'holo_L': 'pdbs/8ir7__AA1_Q7DGD4-L.pdb',\n",
-       "  'predicted_R': 'pdbs/af__P51765.pdb',\n",
-       "  'predicted_L': 'pdbs/af__Q7DGD4.pdb',\n",
+       "({'native': 'pdbs/6iat__C36_Q859I3--6iat__C40_Q859I3.pdb',\n",
+       "  'holo_R': 'pdbs/6iat__C36_Q859I3-R.pdb',\n",
+       "  'holo_L': 'pdbs/6iat__C40_Q859I3-L.pdb',\n",
+       "  'predicted_R': '',\n",
+       "  'predicted_L': '',\n",
        "  'apo_R': '',\n",
        "  'apo_L': '',\n",
        "  'apo_R_alt': [],\n",
        "  'apo_L_alt': []},\n",
-       " {'native': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/8ir7__S1_P51765--8ir7__AA1_Q7DGD4.pdb'),\n",
-       "  'holo_R': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/8ir7__S1_P51765-R.pdb'),\n",
-       "  'holo_L': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/8ir7__AA1_Q7DGD4-L.pdb'),\n",
-       "  'predicted_R': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/af__P51765.pdb'),\n",
-       "  'predicted_L': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/af__Q7DGD4.pdb'),\n",
+       " {'native': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/6iat__C36_Q859I3--6iat__C40_Q859I3.pdb'),\n",
+       "  'holo_R': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/6iat__C36_Q859I3-R.pdb'),\n",
+       "  'holo_L': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/6iat__C40_Q859I3-L.pdb'),\n",
+       "  'predicted_R': '',\n",
+       "  'predicted_L': '',\n",
        "  'apo_R': '',\n",
        "  'apo_L': '',\n",
        "  'apo_R_alt': [],\n",
@@ -309,11 +309,11 @@
     {
      "data": {
       "text/plain": [
-       "{'native': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/6rao__H2_Q6HAD2--6rao__F1_Q6HAD0.pdb'),\n",
-       " 'holo_R_pdb': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/6rao__H2_Q6HAD2-R.pdb'),\n",
-       " 'holo_L_pdb': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/6rao__F1_Q6HAD0-L.pdb'),\n",
-       " 'predicted_R_pdb': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/af__Q6HAD2.pdb'),\n",
-       " 'predicted_L_pdb': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/af__Q6HAD0.pdb'),\n",
+       "{'native': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/7mu1__A37_Q9WZP2--7mu1__A48_Q9WZP2.pdb'),\n",
+       " 'holo_R_pdb': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/7mu1__A37_Q9WZP2-R.pdb'),\n",
+       " 'holo_L_pdb': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/7mu1__A48_Q9WZP2-L.pdb'),\n",
+       " 'predicted_R_pdb': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/af__Q9WZP2.pdb'),\n",
+       " 'predicted_L_pdb': PosixPath('/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/af__Q9WZP2.pdb'),\n",
        " 'apo_R_pdb': '',\n",
        " 'apo_L_pdb': '',\n",
        " 'apo_R_pdbs': [''],\n",
@@ -359,7 +359,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "id": "575049d7-e550-4b59-98db-0d7b045e391a",
    "metadata": {},
    "outputs": [
@@ -367,23 +367,22 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2024-09-05 16:50:41,522 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=7, items=7\n",
-      "2024-09-05 16:50:41,943 | pinder.core.utils.cloud.process_many:23 | INFO : runtime succeeded:      0.42s\n"
+      "2024-09-25 16:32:28,261 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=7, items=7\n"
      ]
     },
     {
      "data": {
       "text/plain": [
        "Structure(\n",
-       "    filepath=/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/6rao__H2_Q6HAD2-R.pdb,\n",
+       "    filepath=/Users/danielkovtun/.local/share/pinder/2024-02/pdbs/7mu1__A37_Q9WZP2-R.pdb,\n",
        "    uniprot_map=None,\n",
-       "    pinder_id='6rao__H2_Q6HAD2-R',\n",
-       "    atom_array=<class 'biotite.structure.AtomArray'> with shape (1737,),\n",
+       "    pinder_id='7mu1__A37_Q9WZP2-R',\n",
+       "    atom_array=<class 'biotite.structure.AtomArray'> with shape (2141,),\n",
        "    pdb_engine='fastpdb',\n",
        ")"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -402,26 +401,26 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "id": "a31f5f1d-7fac-4cde-a1f1-e151afeb17e7",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "array([[ 30.417, -28.254,  74.132],\n",
-       "       [ 31.12 , -28.716,  75.321],\n",
-       "       [ 32.127, -29.81 ,  74.98 ],\n",
-       "       [ 32.824, -29.725,  73.966],\n",
-       "       [ 30.125, -29.23 ,  76.361],\n",
-       "       [ 29.491, -30.413,  75.911],\n",
-       "       [ 32.233, -30.777,  75.9  ],\n",
-       "       [ 32.949, -32.052,  75.825],\n",
-       "       [ 34.468, -31.925,  75.898],\n",
-       "       [ 35.158, -32.93 ,  76.097]], dtype=float32)"
+       "array([[150.746, 157.31 , 247.742],\n",
+       "       [150.862, 158.701, 248.154],\n",
+       "       [150.855, 158.698, 249.688],\n",
+       "       [150.926, 157.633, 250.298],\n",
+       "       [149.724, 159.536, 247.537],\n",
+       "       [149.737, 161.039, 247.82 ],\n",
+       "       [148.383, 161.934, 247.037],\n",
+       "       [146.994, 161.391, 248.026],\n",
+       "       [150.77 , 159.87 , 250.317],\n",
+       "       [150.906, 159.979, 251.762]], dtype=float32)"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -432,17 +431,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 7,
    "id": "3c9bceba-819b-4d32-b95a-86239ad6e1bf",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "'SLLERGLSKLTLNAWKDREGKIPAGSMSAMYNPETIQLDYQTRFDTEDTINTASQSNRYVISEPVGLNLTLLFDSQMPGNTTPIETQLAMLKSLCAVDAATGSPYFLRITWGKMRWENKGWFAGRARDLSVTYTLFDRDATPLRATVQLSLVADESFVIQQSLKTQSAPDRALVSVPDLASLPLLALSAGGVLASSVDYLSLAWDNDLDNLDDFQTGDFLRATK'"
+       "'MEFLKRSFAPLTEKQWQEIDNRAREIFKTQLYGRKFVDVEGPYGWEYAAHPLGEVEVLSDENEVVKWGLRKSLPLIELRATFTLDLWELDNLERGKPNVDLSSLEETVRKVAEFEDEVIFRGCEKSGVKGLLSFEERKIECGSTPKDLLEAIVRALSIFSKDGIEGPYTLVINTDRWINFLKEEAGHYPLEKRVEECLRGGKIITTPRIEDALVVSERGGDFKLILGQDLSIGYEDREKDAVRLFITETFTFQVVNPEALILLK'"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -453,17 +452,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 8,
    "id": "0c11d67a-c186-4710-bf4d-3589b1fe483f",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "'>6rao__H2_Q6HAD2-R\\nSLLERGLSKLTLNAWKDREGKIPAGSMSAMYNPETIQLDYQTRFDTEDTINTASQSNRYVISEPVGLNLTLLFDSQMPGNTTPIETQLAMLKSLCAVDAATGSPYFLRITWGKMRWENKGWFAGRARDLSVTYTLFDRDATPLRATVQLSLVADESFVIQQSLKTQSAPDRALVSVPDLASLPLLALSAGGVLASSVDYLSLAWDNDLDNLDDFQTGDFLRATK'"
+       "'>7mu1__A37_Q9WZP2-R\\nMEFLKRSFAPLTEKQWQEIDNRAREIFKTQLYGRKFVDVEGPYGWEYAAHPLGEVEVLSDENEVVKWGLRKSLPLIELRATFTLDLWELDNLERGKPNVDLSSLEETVRKVAEFEDEVIFRGCEKSGVKGLLSFEERKIECGSTPKDLLEAIVRALSIFSKDGIEGPYTLVINTDRWINFLKEEAGHYPLEKRVEECLRGGKIITTPRIEDALVVSERGGDFKLILGQDLSIGYEDREKDAVRLFITETFTFQVVNPEALILLK'"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -474,7 +473,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 9,
    "id": "3fd4c712-625c-4f80-ace4-b4c1a1b8a600",
    "metadata": {},
    "outputs": [],
@@ -516,7 +515,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 10,
    "id": "ee6840bb-f6db-4911-8bef-a900bf4a9c8e",
    "metadata": {},
    "outputs": [
@@ -690,7 +689,7 @@
        ")"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -736,7 +735,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 11,
    "id": "16e07a40-bc85-412d-bd52-858e4ce8d05b",
    "metadata": {},
    "outputs": [
@@ -759,7 +758,7 @@
        " ))"
       ]
      },
-     "execution_count": 12,
+     "execution_count": 11,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -782,7 +781,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 12,
    "id": "764dae70-980e-4cd3-9393-78a26bcf6136",
    "metadata": {},
    "outputs": [
@@ -792,7 +791,7 @@
        "'PLGSMFERFTDRARRVVVLAQEEARMLNHNYIGTEHILLGLIHEGEGVAAKSLESLGISLEGVRSQVEEIIGQGQQAPSGHIPFTPRAKKVLELSLREALQLGHNYIGTEHILLGLIREGEGVAAQVLVKLGAELTRVRQQVIQLLSGY'"
       ]
      },
-     "execution_count": 13,
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -803,7 +802,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 13,
    "id": "b0d94de4-68bb-456f-9ac5-a6c93ad79613",
    "metadata": {},
    "outputs": [
@@ -817,7 +816,7 @@
        "       [-14.842, -18.494, -12.077]], dtype=float32)"
       ]
      },
-     "execution_count": 14,
+     "execution_count": 13,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -836,7 +835,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 14,
    "id": "a2965877-9130-4b0b-a9a0-e3352cd4d68f",
    "metadata": {},
    "outputs": [
@@ -852,7 +851,7 @@
        "])"
       ]
      },
-     "execution_count": 15,
+     "execution_count": 14,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -886,7 +885,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 15,
    "id": "3e71a783-0929-413a-80e4-8a6ca6f46c74",
    "metadata": {},
    "outputs": [
@@ -896,7 +895,7 @@
        "PinderLoader(split=val, monomers=holo, systems=1958)"
       ]
      },
-     "execution_count": 2,
+     "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -955,7 +954,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 16,
    "id": "2fc9dc84-d7ae-4765-82f1-4df75a45fd7f",
    "metadata": {},
    "outputs": [],
@@ -971,7 +970,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 17,
    "id": "4c3200cc-0f42-4737-91a6-a490f2af55ab",
    "metadata": {},
    "outputs": [
@@ -979,7 +978,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "  0%|▉                                                                                                                                                                                                | 9/1958 [00:02<07:35,  4.28it/s]\n"
+      "  0%|█                                                                                                                                                                                                                                               | 9/1958 [00:02<07:39,  4.24it/s]\n"
      ]
     }
    ],
@@ -998,7 +997,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 18,
    "id": "0625b00a-1208-48e6-908d-08de0be3f221",
    "metadata": {},
    "outputs": [
@@ -1008,7 +1007,7 @@
        "10"
       ]
      },
-     "execution_count": 20,
+     "execution_count": 18,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1019,17 +1018,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 19,
    "id": "d9154449-15c6-4095-b3be-6e64721324dc",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "<pinder.core.loader.dataset.PinderDataset at 0x1fdbac640>"
+       "<pinder.core.loader.dataset.PinderDataset at 0x216abf490>"
       ]
      },
-     "execution_count": 3,
+     "execution_count": 19,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1090,6 +1089,7 @@
     "Each of the `target_complex` and `feature_complex` values are dictionaries with structural properties encoded by the `pinder.core.loader.geodata.structure2tensor` function by default:\n",
     "* `atom_coordinates`\n",
     "* `atom_types`\n",
+    "* `element_types`\n",
     "* `chain_ids`\n",
     "* `residue_coordinates`\n",
     "* `residue_types`\n",
@@ -1100,20 +1100,27 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 20,
    "id": "38ec8e3d-ba0f-4ff2-a8d1-44be72e32ad6",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "{'target_complex': {'atom_types': tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
+       "{'target_complex': {'atom_types': tensor([[0.],\n",
+       "          [1.],\n",
+       "          [2.],\n",
+       "          ...,\n",
+       "          [1.],\n",
+       "          [2.],\n",
+       "          [3.]]),\n",
+       "  'element_types': tensor([[3.],\n",
+       "          [0.],\n",
+       "          [0.],\n",
        "          ...,\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [0., 0., 1.,  ..., 0., 0., 0.]]),\n",
+       "          [0.],\n",
+       "          [0.],\n",
+       "          [2.]]),\n",
        "  'residue_types': tensor([[16.],\n",
        "          [16.],\n",
        "          [16.],\n",
@@ -1128,22 +1135,351 @@
        "          [177.7620, 463.8650, 166.9020],\n",
        "          [177.4130, 465.0800, 167.7550],\n",
        "          [176.8000, 464.9490, 168.8150]]),\n",
-       "  'residue_coordinates': tensor([[131.7500, 429.3090, 163.5360],\n",
-       "          [132.6810, 428.2520, 163.1550],\n",
-       "          [133.5150, 428.6750, 161.9500],\n",
-       "          ...,\n",
-       "          [177.7620, 463.8650, 166.9020],\n",
-       "          [177.4130, 465.0800, 167.7550],\n",
-       "          [176.8000, 464.9490, 168.8150]]),\n",
+       "  'residue_coordinates': tensor([[132.6810, 428.2520, 163.1550],\n",
+       "          [133.5560, 429.2490, 159.5910],\n",
+       "          [133.8750, 432.8980, 160.6290],\n",
+       "          [136.1110, 431.8050, 163.5130],\n",
+       "          [138.3420, 429.8920, 161.0830],\n",
+       "          [138.5230, 432.9090, 158.7600],\n",
+       "          [139.4710, 435.1730, 161.6750],\n",
+       "          [142.1200, 432.6310, 162.6740],\n",
+       "          [143.5890, 432.7500, 159.1600],\n",
+       "          [143.6190, 436.5600, 159.1540],\n",
+       "          [145.2600, 436.8040, 162.5830],\n",
+       "          [147.8380, 434.1320, 161.7330],\n",
+       "          [148.7390, 435.8920, 158.4790],\n",
+       "          [149.0640, 439.2630, 160.2240],\n",
+       "          [151.2220, 437.7670, 162.9840],\n",
+       "          [153.4710, 435.8030, 160.6120],\n",
+       "          [153.9420, 438.9180, 158.4710],\n",
+       "          [156.1140, 440.2660, 161.3150],\n",
+       "          [158.3330, 437.1860, 161.4250],\n",
+       "          [161.8090, 436.4120, 160.1100],\n",
+       "          [161.1120, 432.8370, 158.9350],\n",
+       "          [157.4020, 432.0710, 159.3930],\n",
+       "          [156.5290, 428.3870, 159.0310],\n",
+       "          [154.0400, 428.1160, 156.1500],\n",
+       "          [153.3520, 424.6590, 154.7510],\n",
+       "          [152.8370, 424.8210, 150.9870],\n",
+       "          [150.9700, 422.4160, 148.7320],\n",
+       "          [153.5900, 422.6860, 145.9670],\n",
+       "          [157.3050, 423.3410, 145.6600],\n",
+       "          [158.4860, 426.6730, 144.2700],\n",
+       "          [160.7730, 427.0880, 141.2590],\n",
+       "          [163.9630, 429.0530, 141.9600],\n",
+       "          [166.4230, 427.4330, 139.5550],\n",
+       "          [167.6170, 430.8330, 138.3250],\n",
+       "          [170.5180, 431.9210, 140.5320],\n",
+       "          [169.7800, 435.6680, 140.3160],\n",
+       "          [167.6790, 437.8790, 142.6110],\n",
+       "          [167.3350, 435.1910, 145.2890],\n",
+       "          [168.9090, 437.3000, 148.0580],\n",
+       "          [166.1860, 439.9440, 148.5750],\n",
+       "          [164.1010, 438.1300, 151.1890],\n",
+       "          [163.0190, 438.8690, 154.7450],\n",
+       "          [164.6750, 435.6560, 155.9860],\n",
+       "          [168.1170, 436.6550, 154.6330],\n",
+       "          [169.8040, 437.0930, 158.0010],\n",
+       "          [173.2680, 438.4850, 158.5750],\n",
+       "          [174.9270, 440.8490, 156.1460],\n",
+       "          [175.9170, 440.8010, 152.5000],\n",
+       "          [179.6060, 441.0910, 153.4330],\n",
+       "          [180.8000, 439.9830, 156.8710],\n",
+       "          [184.2540, 439.9980, 158.4580],\n",
+       "          [185.5840, 437.2970, 160.7670],\n",
+       "          [186.4140, 439.9000, 163.4480],\n",
+       "          [182.9430, 441.4570, 163.7500],\n",
+       "          [181.4930, 442.0910, 167.2100],\n",
+       "          [178.0780, 440.5360, 167.9060],\n",
+       "          [176.1650, 440.9290, 171.1790],\n",
+       "          [173.0530, 438.9530, 172.1310],\n",
+       "          [170.3990, 440.6990, 174.2170],\n",
+       "          [166.7880, 440.0710, 175.2260],\n",
+       "          [164.4480, 442.1010, 173.0340],\n",
+       "          [160.9780, 443.6110, 173.3810],\n",
+       "          [158.8330, 443.9000, 170.2730],\n",
+       "          [160.1820, 442.9360, 166.8360],\n",
+       "          [163.5010, 444.7230, 166.1630],\n",
+       "          [164.6310, 442.4290, 163.3160],\n",
+       "          [166.6330, 444.2560, 160.6240],\n",
+       "          [166.1700, 447.5210, 162.5220],\n",
+       "          [168.1350, 449.9450, 164.6440],\n",
+       "          [168.0040, 449.0470, 168.3240],\n",
+       "          [168.4210, 451.0540, 171.5220],\n",
+       "          [169.0440, 450.0360, 175.1300],\n",
+       "          [166.1650, 450.3680, 177.6100],\n",
+       "          [166.7800, 449.8300, 181.3340],\n",
+       "          [163.8480, 448.9150, 183.5760],\n",
+       "          [163.3500, 443.9770, 187.9750],\n",
+       "          [165.0240, 442.4930, 184.8870],\n",
+       "          [168.1120, 443.0980, 182.7810],\n",
+       "          [168.1280, 445.9210, 180.2050],\n",
+       "          [166.7230, 445.0130, 176.8030],\n",
+       "          [167.0690, 446.1330, 173.1930],\n",
+       "          [164.0790, 447.8390, 171.5780],\n",
+       "          [163.4630, 449.0630, 168.0350],\n",
+       "          [164.4160, 452.7320, 168.0130],\n",
+       "          [166.7460, 455.3190, 166.5790],\n",
+       "          [167.5460, 459.0090, 166.8440],\n",
+       "          [170.2010, 459.9620, 169.3810],\n",
+       "          [169.5150, 456.8220, 171.4570],\n",
+       "          [170.7440, 454.3880, 168.7860],\n",
+       "          [172.8900, 451.6810, 170.3850],\n",
+       "          [173.2420, 449.3000, 167.4490],\n",
+       "          [171.4700, 447.3670, 164.7150],\n",
+       "          [169.7840, 443.9650, 164.9670],\n",
+       "          [170.8880, 441.4120, 162.3680],\n",
+       "          [169.3320, 438.1640, 163.6510],\n",
+       "          [166.7170, 436.9060, 166.1030],\n",
+       "          [166.2830, 433.5290, 167.7900],\n",
+       "          [162.5490, 433.4260, 168.5040],\n",
+       "          [162.6440, 430.3540, 170.7620],\n",
+       "          [164.9490, 432.1120, 173.2350],\n",
+       "          [163.6530, 435.5760, 172.2210],\n",
+       "          [167.2540, 436.7450, 171.8300],\n",
+       "          [168.3710, 439.3650, 169.3060],\n",
+       "          [171.8910, 439.6680, 167.8930],\n",
+       "          [173.1360, 443.2500, 167.5870],\n",
+       "          [176.0990, 444.7950, 165.8030],\n",
+       "          [177.1670, 447.8660, 167.8270],\n",
+       "          [177.4160, 451.2430, 166.1360],\n",
+       "          [181.1790, 451.0810, 166.7480],\n",
+       "          [181.3560, 448.4350, 163.9980],\n",
+       "          [180.6020, 448.8920, 160.3060],\n",
+       "          [177.8440, 446.8990, 158.6170],\n",
+       "          [176.7540, 446.3160, 155.0150],\n",
+       "          [173.3820, 444.6900, 154.4770],\n",
+       "          [169.6430, 444.9660, 153.9760],\n",
+       "          [167.6080, 446.9450, 156.5140],\n",
+       "          [163.9630, 447.9270, 156.8800],\n",
+       "          [163.1520, 451.2160, 155.1430],\n",
+       "          [160.3870, 453.6570, 156.0420],\n",
+       "          [158.9330, 453.6920, 152.5150],\n",
+       "          [159.8050, 453.2460, 148.8360],\n",
+       "          [161.2560, 456.7540, 148.4440],\n",
+       "          [164.9420, 455.8330, 148.8520],\n",
+       "          [167.3780, 455.9500, 145.9260],\n",
+       "          [171.0380, 455.0610, 145.5060],\n",
+       "          [173.5830, 457.5340, 146.9270],\n",
+       "          [171.2360, 458.9430, 149.5740],\n",
+       "          [171.8150, 459.8640, 153.2130],\n",
+       "          [169.5910, 457.9090, 155.6000],\n",
+       "          [168.4960, 458.3960, 159.2100],\n",
+       "          [167.0870, 455.9020, 161.7110],\n",
+       "          [163.6430, 457.0390, 162.8610],\n",
+       "          [161.8700, 456.3550, 166.1660],\n",
+       "          [160.7490, 452.9260, 164.8950],\n",
+       "          [164.2440, 451.8750, 163.8110],\n",
+       "          [163.4490, 452.3040, 160.1090],\n",
+       "          [165.6820, 454.0140, 157.5590],\n",
+       "          [164.2940, 457.2310, 156.0670],\n",
+       "          [165.6810, 459.7290, 153.5780],\n",
+       "          [167.3870, 462.7790, 155.0760],\n",
+       "          [170.2060, 465.8340, 161.8360],\n",
+       "          [172.1620, 462.7680, 160.6890],\n",
+       "          [174.0280, 461.7040, 163.8180],\n",
+       "          [175.2110, 458.5660, 161.9930],\n",
+       "          [176.5410, 458.2410, 158.4450],\n",
+       "          [174.5550, 455.7530, 156.3540],\n",
+       "          [174.3580, 455.6030, 152.5600],\n",
+       "          [171.8590, 453.6260, 150.4910],\n",
+       "          [173.6300, 451.2730, 148.0810],\n",
+       "          [170.4820, 450.1750, 146.2000],\n",
+       "          [166.8520, 451.1980, 145.7810],\n",
+       "          [164.0730, 450.1880, 148.1490],\n",
+       "          [162.6710, 446.7750, 147.1980],\n",
+       "          [159.1420, 445.8080, 148.2240],\n",
+       "          [158.3160, 442.3640, 149.5950],\n",
+       "          [155.0300, 443.1470, 151.4050],\n",
+       "          [152.6800, 446.0960, 151.8500],\n",
+       "          [155.0610, 447.4890, 154.4910],\n",
+       "          [158.2010, 445.3460, 154.0190],\n",
+       "          [160.8090, 447.5340, 152.2870],\n",
+       "          [164.4130, 446.3150, 152.1360],\n",
+       "          [167.3360, 448.5380, 151.1380],\n",
+       "          [171.0310, 447.6630, 150.9500],\n",
+       "          [172.8160, 450.1440, 153.2260],\n",
+       "          [176.3550, 450.4540, 154.5850],\n",
+       "          [176.6310, 452.0780, 158.0240],\n",
+       "          [180.1600, 453.0970, 159.0310],\n",
+       "          [179.9110, 454.9710, 162.3160],\n",
+       "          [178.7370, 458.4030, 163.4610],\n",
+       "          [179.1350, 461.9150, 162.0830],\n",
+       "          [180.3160, 464.9940, 163.9950],\n",
+       "          [195.1700, 454.0210, 194.1160],\n",
+       "          [196.4220, 450.5570, 193.1610],\n",
+       "          [195.0530, 447.5260, 191.3340],\n",
+       "          [195.7530, 445.4320, 194.4410],\n",
+       "          [195.6760, 446.8960, 197.9710],\n",
+       "          [199.1230, 447.9300, 199.1640],\n",
+       "          [198.4070, 447.0140, 202.7840],\n",
+       "          [198.1130, 443.3080, 203.5090],\n",
+       "          [195.0350, 441.7530, 205.1060],\n",
+       "          [194.7040, 440.1720, 208.5300],\n",
+       "          [191.7420, 437.9460, 207.6670],\n",
+       "          [190.9460, 435.0990, 205.3140],\n",
+       "          [188.0870, 435.6760, 202.8880],\n",
+       "          [184.8300, 433.7320, 203.1830],\n",
+       "          [184.1270, 431.8260, 199.9570],\n",
+       "          [180.5450, 430.8500, 200.7460],\n",
+       "          [178.2470, 433.3320, 199.0220],\n",
+       "          [178.2770, 433.9260, 195.2680],\n",
+       "          [176.3170, 437.2030, 195.0280],\n",
+       "          [177.7690, 440.7330, 194.9660],\n",
+       "          [181.3350, 439.3910, 195.1330],\n",
+       "          [182.5820, 441.3460, 192.0940],\n",
+       "          [182.3530, 444.9560, 193.3610],\n",
+       "          [186.0870, 445.2240, 194.0290],\n",
+       "          [188.7230, 447.5570, 192.6190],\n",
+       "          [190.8760, 444.5870, 191.5490],\n",
+       "          [188.2670, 443.5230, 188.9590],\n",
+       "          [190.1360, 444.3010, 185.7360],\n",
+       "          [188.8140, 444.1990, 182.1980],\n",
+       "          [185.1300, 444.1560, 181.3470],\n",
+       "          [182.1290, 442.2030, 182.6060],\n",
+       "          [181.3330, 441.2700, 178.9860],\n",
+       "          [184.3690, 440.6570, 176.7700],\n",
+       "          [184.4240, 439.7520, 173.0740],\n",
+       "          [187.1540, 437.5560, 171.6020],\n",
+       "          [187.5170, 439.9750, 168.6500],\n",
+       "          [188.4180, 443.0500, 170.7240],\n",
+       "          [191.4360, 445.2100, 169.8910],\n",
+       "          [193.7480, 445.6510, 172.8940],\n",
+       "          [196.8200, 447.8990, 173.0280],\n",
+       "          [199.4910, 448.6120, 175.6310],\n",
+       "          [201.2080, 451.9090, 176.3780],\n",
+       "          [203.6110, 453.4740, 178.8680],\n",
+       "          [201.6130, 454.9210, 181.7570],\n",
+       "          [202.0690, 458.3840, 183.2730],\n",
+       "          [200.4970, 458.3690, 186.7180],\n",
+       "          [197.7960, 455.8130, 187.5240],\n",
+       "          [195.1460, 455.8470, 184.7720],\n",
+       "          [193.6160, 452.5740, 185.9790],\n",
+       "          [189.8600, 452.3800, 185.3150],\n",
+       "          [189.9440, 455.8580, 183.7680],\n",
+       "          [189.7400, 457.4310, 180.3410],\n",
+       "          [193.1630, 457.9480, 178.8100],\n",
+       "          [194.6770, 460.5170, 176.4640],\n",
+       "          [197.9720, 460.6430, 174.6010],\n",
+       "          [200.7710, 462.5710, 176.3210],\n",
+       "          [203.7970, 463.7400, 174.3330],\n",
+       "          [207.1150, 464.5820, 175.9770],\n",
+       "          [210.1390, 466.5260, 174.8060],\n",
+       "          [210.8850, 463.9390, 172.1200],\n",
+       "          [213.6440, 461.8050, 173.6150],\n",
+       "          [211.1280, 459.2430, 174.9140],\n",
+       "          [208.0390, 457.7610, 173.2930],\n",
+       "          [204.4770, 458.9110, 173.9310],\n",
+       "          [202.5240, 457.8000, 177.0030],\n",
+       "          [198.9190, 457.7220, 178.2440],\n",
+       "          [197.5730, 459.9040, 181.0500],\n",
+       "          [194.2310, 460.2350, 182.8060],\n",
+       "          [192.1140, 462.6430, 180.7850],\n",
+       "          [189.0250, 463.1720, 178.7000],\n",
+       "          [187.1240, 465.8360, 176.8220],\n",
+       "          [188.2850, 466.3320, 173.2420],\n",
+       "          [191.5580, 464.4520, 173.8600],\n",
+       "          [189.9950, 461.1280, 174.9010],\n",
+       "          [191.8440, 458.2970, 173.1540],\n",
+       "          [190.4670, 455.2690, 174.9900],\n",
+       "          [189.7380, 453.5820, 178.3070],\n",
+       "          [192.3390, 451.8830, 180.5000],\n",
+       "          [191.2020, 448.5340, 181.9050],\n",
+       "          [194.3800, 447.0280, 183.4050],\n",
+       "          [197.8700, 447.9650, 184.5850],\n",
+       "          [200.9840, 445.8270, 184.9660],\n",
+       "          [202.8290, 447.6220, 187.7590],\n",
+       "          [206.2120, 445.9180, 187.3400],\n",
+       "          [206.6040, 447.1310, 183.7460],\n",
+       "          [204.4430, 450.2380, 184.1830],\n",
+       "          [202.2310, 449.1370, 181.2890],\n",
+       "          [198.6660, 450.3790, 180.7720],\n",
+       "          [196.2310, 448.2490, 178.7680],\n",
+       "          [194.0550, 450.4240, 176.5260],\n",
+       "          [190.8490, 449.7090, 174.6070],\n",
+       "          [190.4360, 452.3080, 171.8240],\n",
+       "          [187.3830, 454.5590, 171.7980],\n",
+       "          [186.4670, 452.9730, 168.4480],\n",
+       "          [185.9050, 449.6480, 170.2580],\n",
+       "          [182.8890, 448.8910, 172.4330],\n",
+       "          [183.5680, 448.6030, 176.1630],\n",
+       "          [181.3330, 447.5400, 179.0590],\n",
+       "          [182.7340, 447.5270, 182.5670],\n",
+       "          [183.3280, 449.2260, 185.8920],\n",
+       "          [184.8760, 452.6980, 185.6990],\n",
+       "          [185.7140, 455.4170, 188.2030],\n",
+       "          [182.9420, 457.9530, 188.8300],\n",
+       "          [183.2930, 461.5650, 189.9350],\n",
+       "          [182.7580, 462.1320, 193.6660],\n",
+       "          [180.0620, 464.7490, 193.1030],\n",
+       "          [176.4120, 463.9030, 192.5790],\n",
+       "          [177.1550, 460.1730, 192.6930],\n",
+       "          [173.8180, 459.5980, 194.4390],\n",
+       "          [172.0770, 461.1030, 191.4040],\n",
+       "          [173.6950, 458.5410, 189.0910],\n",
+       "          [171.0060, 456.0190, 188.1270],\n",
+       "          [170.8210, 453.3390, 185.4210],\n",
+       "          [170.0180, 454.8270, 182.0350],\n",
+       "          [171.8690, 458.1080, 182.6590],\n",
+       "          [174.2320, 459.6520, 180.1210],\n",
+       "          [177.8640, 459.9130, 181.2230],\n",
+       "          [180.8780, 461.8410, 179.9490],\n",
+       "          [184.5650, 461.8070, 180.8430],\n",
+       "          [186.0650, 464.7990, 182.6440],\n",
+       "          [189.5780, 466.2740, 182.6170],\n",
+       "          [190.7120, 463.6340, 185.1390],\n",
+       "          [189.2980, 460.6630, 183.2290],\n",
+       "          [186.3560, 460.3000, 185.6240],\n",
+       "          [182.7840, 459.6590, 184.5210],\n",
+       "          [180.2130, 462.3320, 185.3590],\n",
+       "          [176.5320, 462.8410, 184.6140],\n",
+       "          [175.8580, 465.0300, 181.5790],\n",
+       "          [178.1930, 465.7110, 174.9000],\n",
+       "          [180.9150, 464.1190, 172.7670],\n",
+       "          [181.2510, 460.5420, 174.1130],\n",
+       "          [178.4430, 457.9660, 174.2140],\n",
+       "          [178.5420, 456.4790, 177.7180],\n",
+       "          [175.4730, 455.1590, 179.5500],\n",
+       "          [175.4670, 454.0690, 183.1870],\n",
+       "          [174.1210, 450.5430, 183.6830],\n",
+       "          [173.9850, 450.6340, 187.5030],\n",
+       "          [174.3640, 452.9240, 190.5240],\n",
+       "          [177.5140, 454.2200, 192.1820],\n",
+       "          [179.3060, 451.6570, 194.3540],\n",
+       "          [181.4890, 453.2110, 197.0570],\n",
+       "          [184.7390, 451.3110, 197.6330],\n",
+       "          [186.4870, 454.0230, 199.6690],\n",
+       "          [185.5270, 457.3810, 201.1790],\n",
+       "          [186.7910, 459.0050, 197.9610],\n",
+       "          [186.6010, 456.0430, 195.5320],\n",
+       "          [183.3680, 455.3620, 193.6270],\n",
+       "          [182.9110, 452.9430, 190.7280],\n",
+       "          [180.0190, 452.5230, 188.3060],\n",
+       "          [179.0910, 450.0250, 185.5960],\n",
+       "          [178.9850, 451.7840, 182.2280],\n",
+       "          [179.0510, 451.0080, 178.5100],\n",
+       "          [181.0530, 453.1550, 176.0780],\n",
+       "          [180.1320, 452.7170, 172.4060],\n",
+       "          [182.0570, 455.4530, 170.6120],\n",
+       "          [181.8140, 459.1820, 170.0770],\n",
+       "          [178.5990, 461.0580, 169.3130],\n",
+       "          [177.7620, 463.8650, 166.9020]]),\n",
        "  'residue_ids': tensor([  4.,   4.,   4.,  ..., 182., 182., 182.]),\n",
        "  'chain_ids': tensor([0., 0., 0.,  ..., 1., 1., 1.])},\n",
-       " 'feature_complex': {'atom_types': tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
+       " 'feature_complex': {'atom_types': tensor([[0.],\n",
+       "          [1.],\n",
+       "          [2.],\n",
        "          ...,\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [1., 0., 0.,  ..., 0., 0., 0.],\n",
-       "          [0., 0., 1.,  ..., 0., 0., 0.]]),\n",
+       "          [1.],\n",
+       "          [2.],\n",
+       "          [3.]]),\n",
+       "  'element_types': tensor([[3.],\n",
+       "          [0.],\n",
+       "          [0.],\n",
+       "          ...,\n",
+       "          [0.],\n",
+       "          [0.],\n",
+       "          [2.]]),\n",
        "  'residue_types': tensor([[16.],\n",
        "          [16.],\n",
        "          [16.],\n",
@@ -1155,16 +1491,338 @@
        "          [132.6810, 428.2520, 163.1550],\n",
        "          [133.5150, 428.6750, 161.9500],\n",
        "          ...,\n",
-       "          [198.0186, 449.2866, 161.9663],\n",
-       "          [196.8790, 450.1770, 161.4823],\n",
-       "          [195.8437, 450.2865, 162.1400]]),\n",
-       "  'residue_coordinates': tensor([[131.7500, 429.3090, 163.5360],\n",
-       "          [132.6810, 428.2520, 163.1550],\n",
-       "          [133.5150, 428.6750, 161.9500],\n",
-       "          ...,\n",
-       "          [198.0186, 449.2866, 161.9663],\n",
-       "          [196.8790, 450.1770, 161.4823],\n",
-       "          [195.8437, 450.2865, 162.1400]]),\n",
+       "          [186.8651, 447.7225, 152.6899],\n",
+       "          [185.4418, 447.8365, 153.2255],\n",
+       "          [185.1119, 447.2595, 154.2621]]),\n",
+       "  'residue_coordinates': tensor([[132.6810, 428.2520, 163.1550],\n",
+       "          [133.5560, 429.2490, 159.5910],\n",
+       "          [133.8750, 432.8980, 160.6290],\n",
+       "          [136.1110, 431.8050, 163.5130],\n",
+       "          [138.3420, 429.8920, 161.0830],\n",
+       "          [138.5230, 432.9090, 158.7600],\n",
+       "          [139.4710, 435.1730, 161.6750],\n",
+       "          [142.1200, 432.6310, 162.6740],\n",
+       "          [143.5890, 432.7500, 159.1600],\n",
+       "          [143.6190, 436.5600, 159.1540],\n",
+       "          [145.2600, 436.8040, 162.5830],\n",
+       "          [147.8380, 434.1320, 161.7330],\n",
+       "          [148.7390, 435.8920, 158.4790],\n",
+       "          [149.0640, 439.2630, 160.2240],\n",
+       "          [151.2220, 437.7670, 162.9840],\n",
+       "          [153.4710, 435.8030, 160.6120],\n",
+       "          [153.9420, 438.9180, 158.4710],\n",
+       "          [156.1140, 440.2660, 161.3150],\n",
+       "          [158.3330, 437.1860, 161.4250],\n",
+       "          [161.8090, 436.4120, 160.1100],\n",
+       "          [161.1120, 432.8370, 158.9350],\n",
+       "          [157.4020, 432.0710, 159.3930],\n",
+       "          [156.5290, 428.3870, 159.0310],\n",
+       "          [154.0400, 428.1160, 156.1500],\n",
+       "          [153.3520, 424.6590, 154.7510],\n",
+       "          [152.8370, 424.8210, 150.9870],\n",
+       "          [150.9700, 422.4160, 148.7320],\n",
+       "          [153.5900, 422.6860, 145.9670],\n",
+       "          [157.3050, 423.3410, 145.6600],\n",
+       "          [158.4860, 426.6730, 144.2700],\n",
+       "          [160.7730, 427.0880, 141.2590],\n",
+       "          [163.9630, 429.0530, 141.9600],\n",
+       "          [166.4230, 427.4330, 139.5550],\n",
+       "          [167.6170, 430.8330, 138.3250],\n",
+       "          [170.5180, 431.9210, 140.5320],\n",
+       "          [169.7800, 435.6680, 140.3160],\n",
+       "          [167.6790, 437.8790, 142.6110],\n",
+       "          [167.3350, 435.1910, 145.2890],\n",
+       "          [168.9090, 437.3000, 148.0580],\n",
+       "          [166.1860, 439.9440, 148.5750],\n",
+       "          [164.1010, 438.1300, 151.1890],\n",
+       "          [163.0190, 438.8690, 154.7450],\n",
+       "          [164.6750, 435.6560, 155.9860],\n",
+       "          [168.1170, 436.6550, 154.6330],\n",
+       "          [169.8040, 437.0930, 158.0010],\n",
+       "          [173.2680, 438.4850, 158.5750],\n",
+       "          [174.9270, 440.8490, 156.1460],\n",
+       "          [175.9170, 440.8010, 152.5000],\n",
+       "          [179.6060, 441.0910, 153.4330],\n",
+       "          [180.8000, 439.9830, 156.8710],\n",
+       "          [184.2540, 439.9980, 158.4580],\n",
+       "          [185.5840, 437.2970, 160.7670],\n",
+       "          [186.4140, 439.9000, 163.4480],\n",
+       "          [182.9430, 441.4570, 163.7500],\n",
+       "          [181.4930, 442.0910, 167.2100],\n",
+       "          [178.0780, 440.5360, 167.9060],\n",
+       "          [176.1650, 440.9290, 171.1790],\n",
+       "          [173.0530, 438.9530, 172.1310],\n",
+       "          [170.3990, 440.6990, 174.2170],\n",
+       "          [166.7880, 440.0710, 175.2260],\n",
+       "          [164.4480, 442.1010, 173.0340],\n",
+       "          [160.9780, 443.6110, 173.3810],\n",
+       "          [158.8330, 443.9000, 170.2730],\n",
+       "          [160.1820, 442.9360, 166.8360],\n",
+       "          [163.5010, 444.7230, 166.1630],\n",
+       "          [164.6310, 442.4290, 163.3160],\n",
+       "          [166.6330, 444.2560, 160.6240],\n",
+       "          [166.1700, 447.5210, 162.5220],\n",
+       "          [168.1350, 449.9450, 164.6440],\n",
+       "          [168.0040, 449.0470, 168.3240],\n",
+       "          [168.4210, 451.0540, 171.5220],\n",
+       "          [169.0440, 450.0360, 175.1300],\n",
+       "          [166.1650, 450.3680, 177.6100],\n",
+       "          [166.7800, 449.8300, 181.3340],\n",
+       "          [163.8480, 448.9150, 183.5760],\n",
+       "          [163.3500, 443.9770, 187.9750],\n",
+       "          [165.0240, 442.4930, 184.8870],\n",
+       "          [168.1120, 443.0980, 182.7810],\n",
+       "          [168.1280, 445.9210, 180.2050],\n",
+       "          [166.7230, 445.0130, 176.8030],\n",
+       "          [167.0690, 446.1330, 173.1930],\n",
+       "          [164.0790, 447.8390, 171.5780],\n",
+       "          [163.4630, 449.0630, 168.0350],\n",
+       "          [164.4160, 452.7320, 168.0130],\n",
+       "          [166.7460, 455.3190, 166.5790],\n",
+       "          [167.5460, 459.0090, 166.8440],\n",
+       "          [170.2010, 459.9620, 169.3810],\n",
+       "          [169.5150, 456.8220, 171.4570],\n",
+       "          [170.7440, 454.3880, 168.7860],\n",
+       "          [172.8900, 451.6810, 170.3850],\n",
+       "          [173.2420, 449.3000, 167.4490],\n",
+       "          [171.4700, 447.3670, 164.7150],\n",
+       "          [169.7840, 443.9650, 164.9670],\n",
+       "          [170.8880, 441.4120, 162.3680],\n",
+       "          [169.3320, 438.1640, 163.6510],\n",
+       "          [166.7170, 436.9060, 166.1030],\n",
+       "          [166.2830, 433.5290, 167.7900],\n",
+       "          [162.5490, 433.4260, 168.5040],\n",
+       "          [162.6440, 430.3540, 170.7620],\n",
+       "          [164.9490, 432.1120, 173.2350],\n",
+       "          [163.6530, 435.5760, 172.2210],\n",
+       "          [167.2540, 436.7450, 171.8300],\n",
+       "          [168.3710, 439.3650, 169.3060],\n",
+       "          [171.8910, 439.6680, 167.8930],\n",
+       "          [173.1360, 443.2500, 167.5870],\n",
+       "          [176.0990, 444.7950, 165.8030],\n",
+       "          [177.1670, 447.8660, 167.8270],\n",
+       "          [177.4160, 451.2430, 166.1360],\n",
+       "          [181.1790, 451.0810, 166.7480],\n",
+       "          [181.3560, 448.4350, 163.9980],\n",
+       "          [180.6020, 448.8920, 160.3060],\n",
+       "          [177.8440, 446.8990, 158.6170],\n",
+       "          [176.7540, 446.3160, 155.0150],\n",
+       "          [173.3820, 444.6900, 154.4770],\n",
+       "          [169.6430, 444.9660, 153.9760],\n",
+       "          [167.6080, 446.9450, 156.5140],\n",
+       "          [163.9630, 447.9270, 156.8800],\n",
+       "          [163.1520, 451.2160, 155.1430],\n",
+       "          [160.3870, 453.6570, 156.0420],\n",
+       "          [158.9330, 453.6920, 152.5150],\n",
+       "          [159.8050, 453.2460, 148.8360],\n",
+       "          [161.2560, 456.7540, 148.4440],\n",
+       "          [164.9420, 455.8330, 148.8520],\n",
+       "          [167.3780, 455.9500, 145.9260],\n",
+       "          [171.0380, 455.0610, 145.5060],\n",
+       "          [173.5830, 457.5340, 146.9270],\n",
+       "          [171.2360, 458.9430, 149.5740],\n",
+       "          [171.8150, 459.8640, 153.2130],\n",
+       "          [169.5910, 457.9090, 155.6000],\n",
+       "          [168.4960, 458.3960, 159.2100],\n",
+       "          [167.0870, 455.9020, 161.7110],\n",
+       "          [163.6430, 457.0390, 162.8610],\n",
+       "          [161.8700, 456.3550, 166.1660],\n",
+       "          [160.7490, 452.9260, 164.8950],\n",
+       "          [164.2440, 451.8750, 163.8110],\n",
+       "          [163.4490, 452.3040, 160.1090],\n",
+       "          [165.6820, 454.0140, 157.5590],\n",
+       "          [164.2940, 457.2310, 156.0670],\n",
+       "          [165.6810, 459.7290, 153.5780],\n",
+       "          [167.3870, 462.7790, 155.0760],\n",
+       "          [170.2060, 465.8340, 161.8360],\n",
+       "          [172.1620, 462.7680, 160.6890],\n",
+       "          [174.0280, 461.7040, 163.8180],\n",
+       "          [175.2110, 458.5660, 161.9930],\n",
+       "          [176.5410, 458.2410, 158.4450],\n",
+       "          [174.5550, 455.7530, 156.3540],\n",
+       "          [174.3580, 455.6030, 152.5600],\n",
+       "          [171.8590, 453.6260, 150.4910],\n",
+       "          [173.6300, 451.2730, 148.0810],\n",
+       "          [170.4820, 450.1750, 146.2000],\n",
+       "          [166.8520, 451.1980, 145.7810],\n",
+       "          [164.0730, 450.1880, 148.1490],\n",
+       "          [162.6710, 446.7750, 147.1980],\n",
+       "          [159.1420, 445.8080, 148.2240],\n",
+       "          [158.3160, 442.3640, 149.5950],\n",
+       "          [155.0300, 443.1470, 151.4050],\n",
+       "          [152.6800, 446.0960, 151.8500],\n",
+       "          [155.0610, 447.4890, 154.4910],\n",
+       "          [158.2010, 445.3460, 154.0190],\n",
+       "          [160.8090, 447.5340, 152.2870],\n",
+       "          [164.4130, 446.3150, 152.1360],\n",
+       "          [167.3360, 448.5380, 151.1380],\n",
+       "          [171.0310, 447.6630, 150.9500],\n",
+       "          [172.8160, 450.1440, 153.2260],\n",
+       "          [176.3550, 450.4540, 154.5850],\n",
+       "          [176.6310, 452.0780, 158.0240],\n",
+       "          [180.1600, 453.0970, 159.0310],\n",
+       "          [179.9110, 454.9710, 162.3160],\n",
+       "          [178.7370, 458.4030, 163.4610],\n",
+       "          [179.1350, 461.9150, 162.0830],\n",
+       "          [180.3160, 464.9940, 163.9950],\n",
+       "          [195.4340, 461.9335, 182.1035],\n",
+       "          [199.2382, 461.9114, 182.0227],\n",
+       "          [201.9684, 459.5313, 180.8911],\n",
+       "          [203.4010, 459.6168, 184.4221],\n",
+       "          [201.2254, 460.1810, 187.5138],\n",
+       "          [201.1625, 463.8214, 188.5711],\n",
+       "          [200.9260, 462.9902, 192.2737],\n",
+       "          [204.0461, 461.5012, 193.8212],\n",
+       "          [204.0666, 458.1489, 195.6112],\n",
+       "          [204.6119, 457.4490, 199.2917],\n",
+       "          [205.8504, 453.8792, 198.8491],\n",
+       "          [208.7257, 452.0784, 197.1862],\n",
+       "          [207.7980, 449.4805, 194.5819],\n",
+       "          [208.4135, 445.7754, 195.1886],\n",
+       "          [210.6612, 444.3414, 192.4637],\n",
+       "          [210.1697, 440.6754, 193.3148],\n",
+       "          [207.5283, 439.2731, 190.9761],\n",
+       "          [207.8597, 439.3451, 187.1904],\n",
+       "          [204.2639, 438.5888, 186.1246],\n",
+       "          [201.5404, 441.1361, 185.3079],\n",
+       "          [203.9241, 444.0497, 185.9199],\n",
+       "          [203.2558, 445.7530, 182.5636],\n",
+       "          [199.5896, 446.7999, 182.9544],\n",
+       "          [200.4436, 450.4320, 183.6891],\n",
+       "          [199.5210, 453.6368, 181.8842],\n",
+       "          [203.2026, 454.6235, 181.6131],\n",
+       "          [203.8955, 451.7040, 179.2375],\n",
+       "          [204.5535, 453.5907, 175.9981],\n",
+       "          [205.0174, 452.1659, 172.5296],\n",
+       "          [204.0168, 448.6487, 171.5671],\n",
+       "          [204.5046, 445.2215, 173.1233],\n",
+       "          [205.9230, 444.0098, 169.7873],\n",
+       "          [208.0113, 446.5703, 167.8936],\n",
+       "          [209.7071, 446.1658, 164.5108],\n",
+       "          [212.9677, 447.9385, 163.6932],\n",
+       "          [211.5590, 448.9689, 160.2797],\n",
+       "          [208.5749, 450.9334, 161.6230],\n",
+       "          [207.8057, 454.4637, 160.4340],\n",
+       "          [207.4872, 456.9118, 163.3431],\n",
+       "          [206.4332, 460.5624, 163.0761],\n",
+       "          [206.0788, 463.4231, 165.5477],\n",
+       "          [203.4709, 466.1752, 165.5817],\n",
+       "          [202.2743, 469.0646, 167.7365],\n",
+       "          [199.6173, 467.7885, 170.1333],\n",
+       "          [196.2579, 469.4412, 170.8270],\n",
+       "          [194.9515, 468.0990, 174.1181],\n",
+       "          [196.1949, 464.7333, 175.3846],\n",
+       "          [195.9083, 462.1365, 172.5968],\n",
+       "          [198.1073, 459.6493, 174.4645],\n",
+       "          [197.1774, 456.0215, 173.7158],\n",
+       "          [194.3844, 457.2024, 171.4133],\n",
+       "          [193.6677, 457.3959, 167.7107],\n",
+       "          [194.6959, 460.7277, 166.2377],\n",
+       "          [193.3968, 462.9168, 163.4242],\n",
+       "          [194.8157, 465.9834, 161.7135],\n",
+       "          [193.5975, 469.3353, 163.0516],\n",
+       "          [194.0021, 472.4931, 160.9684],\n",
+       "          [193.9675, 475.9664, 162.5033],\n",
+       "          [193.4755, 479.4159, 161.0357],\n",
+       "          [196.7072, 479.1398, 159.0498],\n",
+       "          [199.2380, 481.0833, 161.1038],\n",
+       "          [200.4355, 477.9101, 162.8583],\n",
+       "          [201.1263, 474.4388, 161.5022],\n",
+       "          [198.7341, 471.4980, 161.7173],\n",
+       "          [198.3870, 469.4143, 164.8834],\n",
+       "          [196.9634, 466.0464, 165.9660],\n",
+       "          [193.8728, 465.6269, 168.1384],\n",
+       "          [192.0453, 462.6648, 169.6378],\n",
+       "          [189.6005, 461.3958, 167.0345],\n",
+       "          [188.5610, 458.5793, 164.7638],\n",
+       "          [185.9219, 457.6053, 162.2491],\n",
+       "          [186.6821, 458.7163, 158.6997],\n",
+       "          [189.3552, 461.1924, 159.8627],\n",
+       "          [191.6279, 458.6501, 161.5795],\n",
+       "          [195.2363, 459.3693, 160.6062],\n",
+       "          [197.1180, 457.1342, 163.0350],\n",
+       "          [197.6513, 456.0197, 166.6202],\n",
+       "          [199.5711, 457.9891, 169.2464],\n",
+       "          [201.9258, 455.8545, 171.3396],\n",
+       "          [204.0223, 458.4030, 173.2702],\n",
+       "          [204.0654, 462.0508, 174.3389],\n",
+       "          [206.9754, 464.2809, 175.3246],\n",
+       "          [205.3117, 466.7347, 177.6991],\n",
+       "          [208.0995, 469.3304, 177.8176],\n",
+       "          [207.9500, 469.9574, 174.0591],\n",
+       "          [204.2863, 468.9828, 173.6828],\n",
+       "          [205.2161, 466.4121, 171.0357],\n",
+       "          [203.0047, 463.4515, 170.1065],\n",
+       "          [204.5949, 460.3627, 168.5532],\n",
+       "          [202.3941, 458.9509, 165.7863],\n",
+       "          [202.4135, 455.6142, 163.9592],\n",
+       "          [200.5412, 455.9816, 160.6380],\n",
+       "          [197.4668, 453.8614, 159.9738],\n",
+       "          [199.3793, 452.3290, 157.0461],\n",
+       "          [201.8117, 450.7604, 159.5502],\n",
+       "          [200.9911, 447.7578, 161.7207],\n",
+       "          [200.6235, 448.4547, 165.4404],\n",
+       "          [200.1784, 446.1145, 168.4132],\n",
+       "          [199.8529, 447.5739, 171.8820],\n",
+       "          [197.7348, 448.8400, 174.7466],\n",
+       "          [195.1268, 451.4528, 173.8197],\n",
+       "          [192.3477, 453.2565, 175.6605],\n",
+       "          [188.9569, 451.5234, 175.5765],\n",
+       "          [185.5207, 453.1102, 175.8322],\n",
+       "          [183.9646, 452.9510, 179.3080],\n",
+       "          [180.7992, 451.2673, 178.0513],\n",
+       "          [180.4683, 447.5265, 177.5932],\n",
+       "          [184.0987, 446.9803, 178.5930],\n",
+       "          [183.1029, 443.7189, 180.2918],\n",
+       "          [181.8425, 442.4591, 176.9255],\n",
+       "          [185.2580, 443.0276, 175.3308],\n",
+       "          [186.8817, 439.6114, 174.8690],\n",
+       "          [189.8907, 438.4270, 172.8485],\n",
+       "          [189.0413, 438.0299, 169.1817],\n",
+       "          [186.5213, 440.8977, 169.1051],\n",
+       "          [186.4875, 443.5345, 166.3745],\n",
+       "          [187.2138, 447.0861, 167.5288],\n",
+       "          [186.7567, 450.5166, 165.9644],\n",
+       "          [187.8187, 454.0125, 166.9865],\n",
+       "          [185.1741, 456.5018, 168.1070],\n",
+       "          [185.0113, 460.3016, 167.8790],\n",
+       "          [187.2229, 460.5884, 170.9843],\n",
+       "          [189.9023, 458.1826, 169.7563],\n",
+       "          [188.6959, 455.3909, 172.0528],\n",
+       "          [188.3371, 451.7683, 170.9875],\n",
+       "          [184.8400, 450.2790, 171.0850],\n",
+       "          [183.3118, 446.9541, 170.0985],\n",
+       "          [181.7845, 446.9279, 166.6165],\n",
+       "          [183.4827, 449.0777, 160.0576],\n",
+       "          [186.3409, 451.0179, 158.4578],\n",
+       "          [189.4116, 450.1909, 160.6042],\n",
+       "          [190.7996, 446.6886, 161.1856],\n",
+       "          [191.3849, 446.4283, 164.9391],\n",
+       "          [191.1396, 443.1729, 166.9038],\n",
+       "          [191.2964, 442.9523, 170.6910],\n",
+       "          [193.9518, 440.5240, 171.9330],\n",
+       "          [192.9439, 440.5845, 175.6207],\n",
+       "          [190.2838, 441.8338, 178.0451],\n",
+       "          [189.7749, 445.3017, 179.4824],\n",
+       "          [192.2175, 446.2191, 182.2554],\n",
+       "          [190.9081, 448.9066, 184.6113],\n",
+       "          [193.6014, 451.3525, 185.7371],\n",
+       "          [191.2417, 453.9911, 187.1604],\n",
+       "          [187.5045, 454.2755, 187.8168],\n",
+       "          [187.1856, 455.8765, 184.3634],\n",
+       "          [190.3867, 454.6044, 182.6770],\n",
+       "          [190.3628, 451.2543, 180.8545],\n",
+       "          [193.0865, 449.8934, 178.5754],\n",
+       "          [193.0575, 446.9304, 176.2036],\n",
+       "          [195.6519, 445.1075, 174.1079],\n",
+       "          [194.7845, 445.4580, 170.4237],\n",
+       "          [196.3712, 445.1063, 166.9902],\n",
+       "          [195.6411, 447.6097, 164.2096],\n",
+       "          [196.5773, 446.4442, 160.7041],\n",
+       "          [195.1365, 449.0989, 158.4050],\n",
+       "          [191.7716, 450.0980, 157.0157],\n",
+       "          [189.1553, 447.6691, 155.7140],\n",
+       "          [186.8651, 447.7225, 152.6899]]),\n",
        "  'residue_ids': tensor([  4.,   4.,   4.,  ..., 182., 182., 182.]),\n",
        "  'chain_ids': tensor([0., 0., 0.,  ..., 1., 1., 1.])},\n",
        " 'id': '8phr__X4_UNDEFINED--8phr__W4_UNDEFINED',\n",
@@ -1172,7 +1830,7 @@
        " 'target_id': '8phr__X4_UNDEFINED-R--8phr__W4_UNDEFINED-L'}"
       ]
      },
-     "execution_count": 4,
+     "execution_count": 20,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1184,7 +1842,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 21,
    "id": "edd68e94-0782-47a5-a9cb-d6f8d1db00f1",
    "metadata": {},
    "outputs": [
@@ -1192,10 +1850,8 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2024-09-23 15:43:35,437 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=5, items=5\n",
-      "2024-09-23 15:43:35,771 | pinder.core.utils.cloud.process_many:23 | INFO : runtime succeeded:      0.33s\n",
-      "2024-09-23 15:43:35,997 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=7, items=7\n",
-      "2024-09-23 15:43:36,256 | pinder.core.utils.cloud.process_many:23 | INFO : runtime succeeded:      0.26s\n"
+      "2024-09-25 16:33:12,882 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=5, items=5\n",
+      "2024-09-25 16:33:13,349 | pinder.core.utils.cloud:375 | INFO : Gsutil process_many=download_to_filename, threads=7, items=7\n"
      ]
     }
    ],
@@ -1232,150 +1888,154 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 22,
    "id": "608249db-4ba3-4391-b8da-79a8550ba974",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "{'target_complex': {'atom_types': tensor([[[ 0.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           [ 1.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           [ 1.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
+       "{'target_complex': {'atom_types': tensor([[[ 0.],\n",
+       "           [ 1.],\n",
+       "           [ 2.],\n",
        "           ...,\n",
-       "           [ 0.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           [ 1.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           [ 1.,  0.,  0.,  ...,  0.,  0.,  0.]],\n",
+       "           [-1.],\n",
+       "           [-1.],\n",
+       "           [-1.]],\n",
        "  \n",
-       "          [[ 0.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           [ 1.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           [ 1.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
+       "          [[ 0.],\n",
+       "           [ 1.],\n",
+       "           [ 2.],\n",
        "           ...,\n",
-       "           [-1., -1., -1.,  ..., -1., -1., -1.],\n",
-       "           [-1., -1., -1.,  ..., -1., -1., -1.],\n",
-       "           [-1., -1., -1.,  ..., -1., -1., -1.]]]),\n",
-       "  'residue_types': tensor([[[15.],\n",
-       "           [15.],\n",
-       "           [15.],\n",
-       "           ...,\n",
-       "           [15.],\n",
-       "           [15.],\n",
-       "           [15.]],\n",
-       "  \n",
-       "          [[12.],\n",
-       "           [12.],\n",
-       "           [12.],\n",
+       "           [ 1.],\n",
+       "           [ 2.],\n",
+       "           [ 3.]]]),\n",
+       "  'residue_types': tensor([[[ 2.],\n",
+       "           [ 2.],\n",
+       "           [ 2.],\n",
        "           ...,\n",
        "           [-1.],\n",
        "           [-1.],\n",
-       "           [-1.]]]),\n",
-       "  'atom_coordinates': tensor([[[ 194.5170,  179.2660,  -91.8800],\n",
-       "           [ 194.9840,  179.8980,  -90.6150],\n",
-       "           [ 193.8100,  180.4120,  -89.7870],\n",
-       "           ...,\n",
-       "           [ 167.9860,  256.5550,  -26.2400],\n",
-       "           [ 168.9990,  256.8070,  -25.2210],\n",
-       "           [ 170.4040,  256.5630,  -25.7840]],\n",
+       "           [-1.]],\n",
        "  \n",
-       "          [[ -90.1810,  -82.1340, -226.8990],\n",
-       "           [ -90.2500,  -81.9540, -225.4530],\n",
-       "           [ -90.5700,  -80.5080, -225.0850],\n",
+       "          [[13.],\n",
+       "           [13.],\n",
+       "           [13.],\n",
+       "           ...,\n",
+       "           [13.],\n",
+       "           [13.],\n",
+       "           [13.]]]),\n",
+       "  'atom_coordinates': tensor([[[ 289.1240,  316.4350,  305.0190],\n",
+       "           [ 290.4980,  316.8960,  304.8630],\n",
+       "           [ 291.3100,  316.6540,  306.1310],\n",
        "           ...,\n",
        "           [-100.0000, -100.0000, -100.0000],\n",
        "           [-100.0000, -100.0000, -100.0000],\n",
-       "           [-100.0000, -100.0000, -100.0000]]]),\n",
-       "  'residue_coordinates': tensor([[[ 194.5170,  179.2660,  -91.8800],\n",
-       "           [ 194.9840,  179.8980,  -90.6150],\n",
-       "           [ 193.8100,  180.4120,  -89.7870],\n",
-       "           ...,\n",
-       "           [ 167.9860,  256.5550,  -26.2400],\n",
-       "           [ 168.9990,  256.8070,  -25.2210],\n",
-       "           [ 170.4040,  256.5630,  -25.7840]],\n",
+       "           [-100.0000, -100.0000, -100.0000]],\n",
        "  \n",
-       "          [[ -90.1810,  -82.1340, -226.8990],\n",
-       "           [ -90.2500,  -81.9540, -225.4530],\n",
-       "           [ -90.5700,  -80.5080, -225.0850],\n",
+       "          [[ 169.2570,  175.1550,  105.7850],\n",
+       "           [ 170.4880,  175.1580,  106.5660],\n",
+       "           [ 170.5260,  176.3570,  107.5090],\n",
+       "           ...,\n",
+       "           [ 158.9320,  185.6340,  166.0330],\n",
+       "           [ 158.8780,  185.2340,  164.5630],\n",
+       "           [ 158.0010,  184.4790,  164.1430]]]),\n",
+       "  'residue_coordinates': tensor([[[ 290.4980,  316.8960,  304.8630],\n",
+       "           [ 291.8970,  317.5740,  308.3580],\n",
+       "           [ 295.3220,  319.0930,  308.9830],\n",
        "           ...,\n",
        "           [-100.0000, -100.0000, -100.0000],\n",
        "           [-100.0000, -100.0000, -100.0000],\n",
-       "           [-100.0000, -100.0000, -100.0000]]]),\n",
-       "  'residue_ids': tensor([[  1.,   1.,   1.,  ..., 280., 280., 280.],\n",
-       "          [  1.,   1.,   1.,  ..., -99., -99., -99.]]),\n",
-       "  'chain_ids': tensor([[ 0.,  0.,  0.,  ...,  1.,  1.,  1.],\n",
-       "          [ 0.,  0.,  0.,  ..., -1., -1., -1.]])},\n",
-       " 'feature_complex': {'atom_types': tensor([[[ 0.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           [ 1.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           [ 1.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           ...,\n",
-       "           [ 0.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           [ 1.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           [ 1.,  0.,  0.,  ...,  0.,  0.,  0.]],\n",
+       "           [-100.0000, -100.0000, -100.0000]],\n",
        "  \n",
-       "          [[ 0.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           [ 1.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
-       "           [ 1.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
+       "          [[ 170.4880,  175.1580,  106.5660],\n",
+       "           [ 171.8100,  177.5960,  109.1870],\n",
+       "           [ 174.0830,  180.6050,  108.7410],\n",
        "           ...,\n",
-       "           [-1., -1., -1.,  ..., -1., -1., -1.],\n",
-       "           [-1., -1., -1.,  ..., -1., -1., -1.],\n",
-       "           [-1., -1., -1.,  ..., -1., -1., -1.]]]),\n",
-       "  'residue_types': tensor([[[15.],\n",
-       "           [15.],\n",
-       "           [15.],\n",
+       "           [ 153.4260,  188.1030,  169.2780],\n",
+       "           [ 155.3650,  185.4840,  167.3250],\n",
+       "           [ 158.9320,  185.6340,  166.0330]]]),\n",
+       "  'residue_ids': tensor([[ 1.0000e+00,  1.0000e+00,  1.0000e+00,  ..., -9.9000e+01,\n",
+       "           -9.9000e+01, -9.9000e+01],\n",
+       "          [ 2.1500e+02,  2.1500e+02,  2.1500e+02,  ...,  1.1670e+03,\n",
+       "            1.1670e+03,  1.1670e+03]]),\n",
+       "  'chain_ids': tensor([[ 0.,  0.,  0.,  ..., -1., -1., -1.],\n",
+       "          [ 0.,  0.,  0.,  ...,  1.,  1.,  1.]])},\n",
+       " 'feature_complex': {'atom_types': tensor([[[ 0.],\n",
+       "           [ 1.],\n",
+       "           [ 2.],\n",
        "           ...,\n",
-       "           [15.],\n",
-       "           [15.],\n",
-       "           [15.]],\n",
+       "           [-1.],\n",
+       "           [-1.],\n",
+       "           [-1.]],\n",
        "  \n",
-       "          [[12.],\n",
-       "           [12.],\n",
-       "           [12.],\n",
+       "          [[ 0.],\n",
+       "           [ 1.],\n",
+       "           [ 2.],\n",
+       "           ...,\n",
+       "           [ 1.],\n",
+       "           [ 2.],\n",
+       "           [ 3.]]]),\n",
+       "  'residue_types': tensor([[[ 2.],\n",
+       "           [ 2.],\n",
+       "           [ 2.],\n",
        "           ...,\n",
        "           [-1.],\n",
        "           [-1.],\n",
-       "           [-1.]]]),\n",
-       "  'atom_coordinates': tensor([[[ 194.5170,  179.2660,  -91.8800],\n",
-       "           [ 194.9840,  179.8980,  -90.6150],\n",
-       "           [ 193.8100,  180.4120,  -89.7870],\n",
-       "           ...,\n",
-       "           [ 178.4873,  215.5139,  -38.6835],\n",
-       "           [ 178.2736,  215.6114,  -40.1233],\n",
-       "           [ 179.1470,  216.7183,  -40.7253]],\n",
+       "           [-1.]],\n",
        "  \n",
-       "          [[ -90.1810,  -82.1340, -226.8990],\n",
-       "           [ -90.2500,  -81.9540, -225.4530],\n",
-       "           [ -90.5700,  -80.5080, -225.0850],\n",
+       "          [[13.],\n",
+       "           [13.],\n",
+       "           [13.],\n",
+       "           ...,\n",
+       "           [13.],\n",
+       "           [13.],\n",
+       "           [13.]]]),\n",
+       "  'atom_coordinates': tensor([[[ 289.1240,  316.4350,  305.0190],\n",
+       "           [ 290.4980,  316.8960,  304.8630],\n",
+       "           [ 291.3100,  316.6540,  306.1310],\n",
        "           ...,\n",
        "           [-100.0000, -100.0000, -100.0000],\n",
        "           [-100.0000, -100.0000, -100.0000],\n",
-       "           [-100.0000, -100.0000, -100.0000]]]),\n",
-       "  'residue_coordinates': tensor([[[ 194.5170,  179.2660,  -91.8800],\n",
-       "           [ 194.9840,  179.8980,  -90.6150],\n",
-       "           [ 193.8100,  180.4120,  -89.7870],\n",
-       "           ...,\n",
-       "           [ 178.4873,  215.5139,  -38.6835],\n",
-       "           [ 178.2736,  215.6114,  -40.1233],\n",
-       "           [ 179.1470,  216.7183,  -40.7253]],\n",
+       "           [-100.0000, -100.0000, -100.0000]],\n",
        "  \n",
-       "          [[ -90.1810,  -82.1340, -226.8990],\n",
-       "           [ -90.2500,  -81.9540, -225.4530],\n",
-       "           [ -90.5700,  -80.5080, -225.0850],\n",
+       "          [[ 169.2570,  175.1550,  105.7850],\n",
+       "           [ 170.4880,  175.1580,  106.5660],\n",
+       "           [ 170.5260,  176.3570,  107.5090],\n",
+       "           ...,\n",
+       "           [ 109.5580,  141.3806,  149.9941],\n",
+       "           [ 110.3494,  140.1194,  150.3209],\n",
+       "           [ 111.2073,  140.1200,  151.2038]]]),\n",
+       "  'residue_coordinates': tensor([[[ 290.4980,  316.8960,  304.8630],\n",
+       "           [ 291.8970,  317.5740,  308.3580],\n",
+       "           [ 295.3220,  319.0930,  308.9830],\n",
        "           ...,\n",
        "           [-100.0000, -100.0000, -100.0000],\n",
        "           [-100.0000, -100.0000, -100.0000],\n",
-       "           [-100.0000, -100.0000, -100.0000]]]),\n",
-       "  'residue_ids': tensor([[  1.,   1.,   1.,  ..., 280., 280., 280.],\n",
-       "          [  1.,   1.,   1.,  ..., -99., -99., -99.]]),\n",
-       "  'chain_ids': tensor([[ 0.,  0.,  0.,  ...,  1.,  1.,  1.],\n",
-       "          [ 0.,  0.,  0.,  ..., -1., -1., -1.]])},\n",
-       " 'id': ['2fte__A43_P49861--2fte__G44_P49861',\n",
-       "  '3jc8__JB1_Q306N5--3jc8__NB1_Q306N4'],\n",
-       " 'sample_id': ['2fte__A43_P49861-R--2fte__G44_P49861-L',\n",
-       "  '3jc8__JB1_Q306N5-R--af__Q306N4'],\n",
-       " 'target_id': ['2fte__A43_P49861-R--2fte__G44_P49861-L',\n",
-       "  '3jc8__JB1_Q306N5-R--3jc8__NB1_Q306N4-L']}"
+       "           [-100.0000, -100.0000, -100.0000]],\n",
+       "  \n",
+       "          [[ 170.4880,  175.1580,  106.5660],\n",
+       "           [ 171.8100,  177.5960,  109.1870],\n",
+       "           [ 174.0830,  180.6050,  108.7410],\n",
+       "           ...,\n",
+       "           [ 106.3922,  144.9902,  154.8819],\n",
+       "           [ 109.4120,  143.4231,  153.1913],\n",
+       "           [ 109.5580,  141.3806,  149.9941]]]),\n",
+       "  'residue_ids': tensor([[ 1.0000e+00,  1.0000e+00,  1.0000e+00,  ..., -9.9000e+01,\n",
+       "           -9.9000e+01, -9.9000e+01],\n",
+       "          [ 2.1500e+02,  2.1500e+02,  2.1500e+02,  ...,  1.1670e+03,\n",
+       "            1.1670e+03,  1.1670e+03]]),\n",
+       "  'chain_ids': tensor([[ 0.,  0.,  0.,  ..., -1., -1., -1.],\n",
+       "          [ 0.,  0.,  0.,  ...,  1.,  1.,  1.]])},\n",
+       " 'id': ['7wl3__A23_A0A866W289--7wl3__A24_A0A866W289',\n",
+       "  '8ae1__A1_Q9RRB6--8ae1__B1_Q9RRB6'],\n",
+       " 'sample_id': ['7wl3__A23_A0A866W289-R--7wl3__A24_A0A866W289-L',\n",
+       "  '8ae1__A1_Q9RRB6-R--8ae1__B1_Q9RRB6-L'],\n",
+       " 'target_id': ['7wl3__A23_A0A866W289-R--7wl3__A24_A0A866W289-L',\n",
+       "  '8ae1__A1_Q9RRB6-R--8ae1__B1_Q9RRB6-L']}"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 22,
      "metadata": {},
      "output_type": "execute_result"
     }
diff --git a/src/pinder-core/pinder/core/loader/dataset.py b/src/pinder-core/pinder/core/loader/dataset.py
index f817af3..6a0c89e 100644
--- a/src/pinder-core/pinder/core/loader/dataset.py
+++ b/src/pinder-core/pinder/core/loader/dataset.py
@@ -52,10 +52,12 @@
 
 
 def structure2tensor_transform(structure: Structure) -> dict[str, torch.Tensor]:
+    calpha = structure.filter("atom_name", mask=["CA"])
     props: dict[str, torch.Tensor] = structure2tensor(
         atom_coordinates=structure.coords,
-        atom_types=structure.atom_array.element,
-        residue_coordinates=structure.coords,
+        atom_types=structure.atom_array.atom_name,
+        element_types=structure.atom_array.element,
+        residue_coordinates=calpha.coords,
         residue_types=structure.atom_array.res_name,
         residue_ids=structure.atom_array.res_id,
         chain_ids=structure.atom_array.chain_id,
diff --git a/src/pinder-core/pinder/core/loader/geodata.py b/src/pinder-core/pinder/core/loader/geodata.py
index 5d95aba..8e75fae 100644
--- a/src/pinder-core/pinder/core/loader/geodata.py
+++ b/src/pinder-core/pinder/core/loader/geodata.py
@@ -30,6 +30,7 @@
 def structure2tensor(
     atom_coordinates: NDArray[np.double] | None = None,
     atom_types: NDArray[np.str_] | None = None,
+    element_types: NDArray[np.str_] | None = None,
     residue_coordinates: NDArray[np.double] | None = None,
     residue_ids: NDArray[np.int_] | None = None,
     residue_types: NDArray[np.str_] | None = None,
@@ -38,14 +39,16 @@ def structure2tensor(
 ) -> dict[str, torch.Tensor]:
     property_dict = {}
     if atom_types is not None:
-        types_array_ele = np.zeros(
-            (len(atom_types), len(set(list(pc.ELE2NUM.values()))))
-        )
+        unknown_name_idx = max(pc.ALL_ATOM_POSNS.values()) + 1
+        types_array_at = np.zeros((len(atom_types), 1))
         for i, name in enumerate(atom_types):
-            types_array_ele[i, pc.ELE2NUM.get(name, "C")] = 1.0
-
-        property_dict["atom_types"] = torch.tensor(types_array_ele).type(dtype)
-
+            types_array_at[i] = pc.ALL_ATOM_POSNS.get(name, unknown_name_idx)
+        property_dict["atom_types"] = torch.tensor(types_array_at).type(dtype)
+    if element_types is not None:
+        types_array_ele = np.zeros((len(element_types), 1))
+        for i, name in enumerate(element_types):
+            types_array_ele[i] = pc.ELE2NUM.get(name, pc.ELE2NUM["other"])
+        property_dict["element_types"] = torch.tensor(types_array_ele).type(dtype)
     if residue_types is not None:
         unknown_name_idx = max(pc.AA_TO_INDEX.values()) + 1
         types_array_res = np.zeros((len(residue_types), 1))
@@ -119,14 +122,16 @@ def from_structure_pair(
         lig_calpha = ligand.filter("atom_name", mask=["CA"])
         rec_props = structure2tensor(
             atom_coordinates=receptor.coords,
-            atom_types=receptor.atom_array.element,
+            atom_types=receptor.atom_array.atom_name,
+            element_types=receptor.atom_array.element,
             residue_coordinates=rec_calpha.coords,
             residue_types=rec_calpha.atom_array.res_name,
             residue_ids=rec_calpha.atom_array.res_id,
         )
         lig_props = structure2tensor(
             atom_coordinates=ligand.coords,
-            atom_types=ligand.atom_array.element,
+            atom_types=ligand.atom_array.atom_name,
+            element_types=ligand.atom_array.element,
             residue_coordinates=lig_calpha.coords,
             residue_types=lig_calpha.atom_array.res_name,
             residue_ids=lig_calpha.atom_array.res_id,


Test patch:
diff --git a/tests/core/loader/test_dataset.py b/tests/core/loader/test_dataset.py
index 84e7651..1793b55 100644
--- a/tests/core/loader/test_dataset.py
+++ b/tests/core/loader/test_dataset.py
@@ -20,6 +20,7 @@
 
 EXPECTED_STRUCTURE_KEYS = {
     "atom_types",
+    "element_types",
     "residue_types",
     "atom_coordinates",
     "residue_coordinates",
@@ -31,7 +32,6 @@
 def test_pairedpdb_heterodata(pinder_temp_dir):
     pinder_id = "3s9d__B1_P48551--3s9d__A1_P01563"
     system = PinderSystem(pinder_id)
-
     nodes = {NodeRepresentation("atom"), NodeRepresentation("residue")}
     holo_data = PairedPDB.from_pinder_system(
         system=system,
@@ -53,8 +53,8 @@ def test_pairedpdb_heterodata(pinder_temp_dir):
     expected_num_feats = {
         "ligand_residue": 0,
         "receptor_residue": 0,
-        "ligand_atom": 12,
-        "receptor_atom": 12,
+        "ligand_atom": 1,
+        "receptor_atom": 1,
     }
     for k, v in expected_num_feats.items():
         assert holo_data.num_node_features[k] == v
@@ -75,8 +75,8 @@ def test_pairedpdb_heterodata(pinder_temp_dir):
     expected_num_feats = {
         "ligand_residue": 0,
         "receptor_residue": 0,
-        "ligand_atom": 12,
-        "receptor_atom": 12,
+        "ligand_atom": 1,
+        "receptor_atom": 1,
     }
     for k, v in expected_num_feats.items():
         assert apo_data.num_node_features[k] == v
@@ -127,8 +127,8 @@ def test_ppi_dataset(split, ids, parallel, limit_by, pinder_temp_dir, capfd):
     expected_num_feats = {
         "ligand_residue": 0,
         "receptor_residue": 0,
-        "ligand_atom": 12,
-        "receptor_atom": 12,
+        "ligand_atom": 1,
+        "receptor_atom": 1,
         "pdb": 0,
     }
     for k, v in expected_num_feats.items():