Raul

Make neighbor list compatible with float16 and bfloat16

5

This PR adds overloads for float16 and bfloat16 input/output of the neighbor extension. Makes it possible to train in float16 or bfloat16 without casting in between.

Print warnings when incompatible parameters are used

4

The OP in this issue https://github.com/openmm/openmm-torch/issues/133 was led to believe the Graph Network is equivariant because TorchMD-Net provides the "equivariant_group" parameter (which is only used by TensorNet). This happens a...

Support higher dimension outputs in TorchMD-Net

14

Right now the TorchMD_Net module is tied to the idea that only one thing (and maybe its derivative) is returned: https://github.com/torchmd/torchmd-net/blob/a11684710d6ba22b29acbf709bfbc05b93f3402d/torchmdnet/models/model.py#L300 With TensorNet it might be interesting to have more...

Log epoch real time in LNNP

7

This PR changes the LNNP module so that the real time since the start is logged each epoch. Allows to track training time with the CSVLogger.

Cell list Distance module does not support triclinic domains

4

Current implementation of the cell list will fail if the provided box vectors do not represent a rectangular domain: https://github.com/torchmd/torchmd-net/blob/e20876ff79fe952db25e5fc55cc051a919ef218f/torchmdnet/neighbors/neighbors_cuda_cell.cuh#L351-L354 Supporting triclinic requires modifying the cell traversal code: https://github.com/torchmd/torchmd-net/blob/e20876ff79fe952db25e5fc55cc051a919ef218f/torchmdnet/neighbors/neighbors_cuda_cell.cuh#L331-L332 similarly...

[WIP] Support arbitrary outputs in TorchMD_Net

5

Following the discussion in https://github.com/torchmd/torchmd-net/issues/198 this PR attempts to give TorchMD_Net the ability to return more than one output ("y") and its derivative ("neg_dy"). This PR is still a draft...

This PR makes it possible to get tensor outputs from TensorNet. Also, it implements some tests to ensure equivariance. TODO: - [ ] Write a somewhat general tensor output model

Mamba fails to remove cudatoolkit

2

### Troubleshooting docs - [X] My problem is not solved in the Troubleshooting docs ### Anaconda default channels - [X] I do NOT use the Anaconda default channels (pkgs/* etc.)...

Update gemm_api_3x.md

1

Fixed what it seems to be an obvious typo.

Train only an output model, freezing the representation model.

1

Adds the `--freeze-representation`, `--reset-output-model` and `--overwrite-representation` to train.py. - Freeze representation: Makes it so that the representation model weights are not trained - Reset output model: Makes it so that...