I am trying to perform certain operations on a single image, while in a training loop.
In case of
batch_size = 1 , it could be easily done by using
torch.squeeze but I am unable to think of a way when I can do it for other batch sizes.
Below is the minimum code for representation -
def train(n_epochs, loaders, model, optimizer, criterion, use_cuda, save_path): for epoch in range(1, n_epochs+1): for batch_idx, (data, target) in enumerate(final_train_loader): # Here the target shape would be B*H*W*N
B: Batch size H,W,N:Height,width,no. of channels
I want it in a form of
H * W * N