Fix prediction code to make it work with tensors other than elastic #47

	name: Testing

	on: [push, pull_request]

	jobs:
	lint:
	strategy:
	max-parallel: 3
	matrix:
	os: [ubuntu-latest]
	python-version: ["3.10"]
	runs-on: ${{ matrix.os }}
	steps:
	- uses: actions/[email protected]

	- name: Set up Python ${{ matrix.python-version }}
	uses: actions/setup-python@v4
	with:
	python-version: ${{ matrix.python-version }}
	cache: pip
	cache-dependency-path: pyproject.toml

	- name: Install dependencies
	run: \|
	python -m pip install --upgrade pip
	pip install pre-commit

	- name: Lint
	run: pre-commit run --all-files --show-diff-on-failure

	test:
	strategy:
	max-parallel: 3
	matrix:
	os: [ubuntu-latest]
	python-version: ["3.8", "3.9", "3.10"]
	runs-on: ${{ matrix.os }}
	steps:
	- uses: actions/checkout@v3

	- name: Set up Python ${{ matrix.python-version }}
	uses: actions/setup-python@v4
	with:
	python-version: ${{ matrix.python-version }}
	cache: pip
	cache-dependency-path: pyproject.toml

	- name: Install package
	run: \|
	python -m pip install --upgrade pip
	pip install torch==2.1.1
	pip install torch-scatter -f https://data.pyg.org/whl/torch-2.1.0+cpu.html
	pip install -e .[strict,test]

	- name: Test
	run: pytest

Provide feedback