Yes, this is generally the right approach.
Could you check, if all the data and target tensors have the same shape, and thus a tensor creation via:
features = [tensor1, tensor2, ...]
features = torch.stack(features)

would work?
I guess the new error is raised because of unexpected shapes, but am unsure which operation raises it.