From 62d58f91322da1e3c7cc93841ecb778f7a074aa9 Mon Sep 17 00:00:00 2001 From: Ian Slagle Date: Tue, 16 Jan 2024 20:11:33 -0500 Subject: [PATCH] Instantiate natoms with device --- matdeeplearn/preprocessor/helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/matdeeplearn/preprocessor/helpers.py b/matdeeplearn/preprocessor/helpers.py index b3cf157c..9d47b13e 100644 --- a/matdeeplearn/preprocessor/helpers.py +++ b/matdeeplearn/preprocessor/helpers.py @@ -72,7 +72,7 @@ def calculate_edges_master( cell = cell.view(1, 3, 3) edge_index, cell_offsets, neighbors = radius_graph_pbc( - r, n_neighbors, pos, cell, torch.tensor([len(pos)]), [True, True, True], offset_number + r, n_neighbors, pos, cell, torch.tensor([len(pos)], device = device), [True, True, True], offset_number ) ocp_out = get_pbc_distances(