Repository: pykeen/pykeen
Issue: OOM Crash on MPS/Apple silicon
Root cause: External Dependency Breakage

### Describe the bug

All models crash (OOM error) in evaluation phase, when using "device="mps"".  i.e., use of Apple silicon accelerator.
They run fine without the mps (cpu on same device). Disabling/enavling earlyt stopping etc' doesn't affect it. It happens across diff models including the nations example. Training works ok, but it crashes in (Default) evaluation. 

### How to reproduce

Run on a device with apple silicon (e.g. M1-M3 macs). Latest library versions are installed. 

```python
from pykeen.datasets import get_dataset
from pykeen.models import ERModel, RotatE, TransR
from pykeen.nn.init import LabelBasedInitializer
from pykeen.pipeline import pipeline

dataset = get_dataset(dataset="nations", dataset_kwargs=dict(create_inverse_triples=True))

# relevant?
# entity_initializer = LabelBasedInitializer.from_triples_factory(
#    triples_factory=dataset.training,
#    for_entities=True,
#)

results = pipeline(
    dataset=dataset,
    model="distmult",
    epochs=11,
    evaluation_fallback=True,
    stopper='early',
    device="mps",  # if this line is disabled (i.e run on cpu), then completes successfully
)
```

### Environment

| Key             | Value                    |
|-----------------|--------------------------|
| OS              | posix                    |
| Platform        | Darwin                   |
| Release         | 23.4.0                   |
| Time            | Wed May  8 12:43:17 2024 |
| Python          | 3.9.19                   |
| PyKEEN          | 1.10.2                   |
| PyKEEN Hash     | UNHASHED                 |
| PyKEEN Branch   |                          |
| PyTorch         | 2.2.2                    |
| CUDA Available? | false                    |
| CUDA Version    | N/A                      |
| cuDNN Version   | N/A                      |`

### Additional information

_No response_

### Issue Template Checks

- [X] This is not a feature request (use a different issue template if it is)
- [X] This is not a question (use the [discussions](https://github.com/pykeen/pykeen/discussions/new) forum instead)
- [X] I've read the text explaining why including environment information is important and understand if I omit this information that my issue will be dismissed

PyKEEN does not yet support MPS yet for a few reasons. This is being worked on in https://github.com/pykeen/pykeen/pull/975.

When giving an issue, it would be helpful to be more specific about what code you ran and exactly what messages were returned.

Running the following works on my system:

```python
results = pipeline(
    dataset=dataset,
    model="distmult",
    epochs=11,
    evaluation_fallback=True,
    stopper='early',
    device="mps",
)
```

Here is my configuration (M3 Mac):

| Key             | Value                    |
|-----------------|--------------------------|
| OS              | posix                    |
| Platform        | Darwin                   |
| Release         | 23.3.0                   |
| Time            | Tue May 14 12:10:39 2024 |
| Python          | 3.12.3                   |
| PyKEEN          | 1.10.2                   |
| PyKEEN Hash     | UNHASHED                 |
| PyKEEN Branch   |                          |
| PyTorch         | 2.3.0                    |
| CUDA Available? | false                    |
| CUDA Version    | N/A                      |
| cuDNN Version   | N/A                      |

Running the following does not work on my system:

```python
results = pipeline(
    dataset=dataset,
    model="distmult",
    epochs=11,
    evaluation_fallback=False,  # this changed from True to False
    stopper='early',
    device="mps",
)
```

It gives the following message:

```python-traceback
WARNING:torch_max_mem.api:Encountered tensors on device_types={'mps'} while only ['cuda'] are considered safe for automatic memory utilization maximization. This may lead to undocumented crashes (but can be safe, too).
                                                                                                                                                                                     /AppleInternal/Library/BuildRoots/4e1473ee-9f66-11ee-8daf-cedaeb4cabe2/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSCore/Types/MPSNDArray.mm:120: failed assertion `[MPSNDArrayDescriptor sliceDimension:withSubrange:] error: subRange.start (199) is not less than length of dimension[0] (3)'
fish: Job 1, 'python' terminated by signal SIGABRT (Abort)
/opt/homebrew/Cellar/python@3.12/3.12.3/Frameworks/Python.framework/Versions/3.12/lib/python3.12/multiprocessing/resource_tracker.py:254: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
```

This is also dependent on https://github.com/mberr/torch-max-mem/issues/14. I'm going to close this issue in favor of https://github.com/pykeen/pykeen/pull/975.


Hi @ddofer ,

if you want to help with improving MPS support, I would be very happy to have someone with access to mps hardware to work on https://github.com/mberr/torch-max-mem/issues/14

Regarding this issue, you can either

* use `device="cpu"`; this might be considerably slower
* explicitly provide an evaluation batch size
```python
results = pipeline(
    evaluation_kwargs=dict(batch_size=256),
    device="mps",
    ...,
)
```
you may need to fiddle around with this batch size to have a reasonable speed while still not facing any OOM issues.