Hi Hell!



 Hell_Razor:

The model is working perfectly


From what I understand calling your model in your training loop
works, but it fails in your test loop.  The error message you quote
suggests that you are applying your model to an input tensor of
the wrong shape.


            # print("Test: ", data.location.size)



As an aside, this looks fishy.  It’s only a comment, of course, but
if the data you pass to your model really has a location property,
it probably isn’t what your model is expecting.


            output = model(data)



I assume that your training loop also has an equivalent
model (data) call in it somewhere.
Print out the shape of data just before your training-loop
model (data) call and your test-loop model (data) call.
I expect that they will not be the same.
If they are the same, capture the training-loop and test-loop
data tensors, and post a complete, self-contained, runnable
script that reproduces your error that uses those example
data tensors.
That is, get rid of the data loaders and loops – just have
something like:
train_data = torch.tensor (...)
model (train_data)
test_data = torch.tensor (...)
model.eval()
with torch.no_grad():
    model (test_data)

And also post the complete results of your self-contained,
runnable script.
Best.
K. Frank