folivo/PINA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Additional improvement related to #395

Browse Source
This commit is contained in:
FilippoOlivo
2025-01-24 09:47:57 +01:00
committed by Nicola Demo
parent afb1bca245
commit 629a6ee43b
2 changed files with 22 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
pina/label_tensor.py
View File

@@ -380,20 +380,21 @@ class LabelTensor(torch.Tensor):
"""
old_dof = old_labels[to_update_dim]['dof']
label_name = old_labels[dim]['name']
# Handle slicing
if isinstance(index, slice):
# Handle slicing
to_update_labels[dim] = {'dof': old_dof[index], 'name': label_name}
# Handle single integer index
elif isinstance(index, int):
# Handle single integer index
to_update_labels[dim] = {'dof': [old_dof[index]],
'name': label_name}
# Handle lists or tensors
elif isinstance(index, (list, torch.Tensor)):
# Handle lists or tensors
indices = [index] if isinstance(index, (int, str)) else index
# Handle list of bools
if isinstance(index, torch.Tensor) and index.dtype == torch.bool:
index = index.nonzero().squeeze()
to_update_labels[dim] = {
'dof': [old_dof[i] for i in indices] if isinstance(old_dof,
list) else indices,
'dof': [old_dof[i] for i in index] if isinstance(old_dof,
list) else index,
'name': label_name
}
else: