From 901e45b2178ff39b4725581ac5b4716e3696f470 Mon Sep 17 00:00:00 2001 From: Krishna Kumar Date: Mon, 15 Jul 2024 09:40:59 -0600 Subject: [PATCH] Fix prepare data --- gns/train.py | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/gns/train.py b/gns/train.py index b498d23..1ce329d 100644 --- a/gns/train.py +++ b/gns/train.py @@ -413,24 +413,26 @@ def setup_tensorboard(cfg, metadata): def prepare_data(example, device_id): - """Prepare data for training or validation.""" - position = example[0][0].to(device_id) - particle_type = example[0][1].to(device_id) - - if len(example[0]) == 4: # if data loader includes material_property - material_property = example[0][2].to(device_id) - n_particles_per_example = example[0][3].to(device_id) - elif len(example[0]) == 3: - material_property = None - n_particles_per_example = example[0][2].to(device_id) + features, labels = example + + if len(features) == 4: # If material property is present + position, particle_type, material_property, n_particles_per_example = features else: - raise ValueError("Unexpected number of elements in the data loader") - - labels = example[1].to(device_id) - + position, particle_type, n_particles_per_example = features + material_property = None + + # Convert numpy arrays to tensors + position = torch.from_numpy(position).float().to(device_id) + particle_type = torch.from_numpy(particle_type).long().to(device_id) + if material_property is not None: + material_property = torch.from_numpy(np.array(material_property)).float().to(device_id) + n_particles_per_example = torch.tensor([n_particles_per_example], device=device_id).long() + labels = torch.from_numpy(labels).float().to(device_id) + return position, particle_type, material_property, n_particles_per_example, labels + def train(rank, cfg, world_size, device, verbose, use_dist): """Train the model.