The code is a bit hard to follow in screenshots and it’s a better idea to post the code directly by wrapping it in three backticks ``` 
However, based on the screenshots, I assume the error is raised in torch.cat((out1, out2), 0)?
Could you explain the indexing a bit?
It seems you are slicing input in the batch dimension to create out2 and I’m not sure why you would want to remove some samples.
Could you post an executable code snippet with some randomly initialized inputs, so that we could debug the code?