Create ptr tensor on same device as Z when batching atomic graphs

[thomasloux]

Error:
ptr (pointer to first atom indices for each system in a batched AtomicGraph) is always initialized on cpu, even if the input sequence is on another device.

Reproduce:

from graph_pes.atomic_graph import AtomicGraph, to_batch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
atoms = bulk("Si", "diamond", a=5.43, cubic=True)
atoms = AtomicGraph.from_ase(atoms).to(device)
print(atoms.Z.device) # cuda:0
batch = to_batch([atoms])
print(batch.Z.device, batch.ptr.device) #cuda:0 cpu

Solve by creating the ptr tensor on same device as Z.