Facing IndexOutOfBounds Error when trying to do prediction

[RajeshSiniya]

Error:

patch shape : torch.Size([784, 1536])
I saved the memory bank and loaded with below code
memory_bank = torch.load('memory_bank.pt')
memory_bank shape : torch.Size([1, 3920000, 1536])

In Predict() function:
s_idx: 2451854
patch shape: torch.Size([784, 1536])
0%| | 0/1 [00:15<?, ?it/s]
Traceback (most recent call last):
File "/home/jupyter/ml_dev/multiclass_anomaly_coreset.py", line 532, in
image_rocauc, pixel_rocauc = patch_core.evaluate(test_dl)
File "/home/jupyter/ml_dev/multiclass_anomaly_coreset.py", line 389, in evaluate
score, segm_map = self.predict(sample) # Anomaly Detection
File "/home/jupyter/ml_dev/multiclass_anomaly_coreset.py", line 437, in predict
m_star = self.memory_bank[dist_score_idxs[s_idx]].unsqueeze(0) # Memory bank patch closest neighbour to m_test_star
IndexError: index 2451854 is out of bounds for dimension 0 with size 1

I made below Current Fix however would like to know your opinion on this
self.memory_bank = self.memory_bank.squeeze(0)
print("memory_bank shape : ", self.memory_bank.shape)
memory_bank shape : torch.Size([3920000, 1536])