Thanx ptrblck for your reply. I am really stuck after spending a couple of days on this problem.
What I want to do:
I want to find the average angular variance of head classes( say classes with samples greater than 500 are head classes, others are tail classes)  and then transform it to tail classes. I think the solution proposed by you in the last post is fine. So initially I want to extract the features from all classes then associate those features to the corresponding classes and then find the angles between features and corresponding class center. The solution that you provided also makes sense to me, however, in my case feature size is 128,2,128 instead of 4. Could you please, help with this, I shall be very thankful to you.
The corresponding code -snippet is given below.
def train(train_loader, model, criterion, optimizer, epoch, opt, train_targets):
“”“one epoch training”""
model.train()

batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()

end = time.time()
for idx, (images, labels) in enumerate(train_loader):
    data_time.update(time.time() - end)

    images = torch.cat([images[0].unsqueeze(1), images[1].unsqueeze(1)],
                       dim=1)
    
    images = images.view(-1, 3, 32, 32).cuda(non_blocking=True)
   # print('..............Image Size...........',images.shape)
    labels = labels.cuda(non_blocking=True)
   # print('..............Label Size...........',labels.shape)
    bsz = labels.shape[0]
         
  
    # compute loss
    features = model(images)
    #print('....................Size of Features (Before)............',features.shape)
       
    features = features.view(bsz, 2, -1)     # feature size 128,2,128
   
    # For Angular variance 
    if opt.dataset == 'cifar10':
        nbclasses=10
        print('Dataset :  Cifar10')
    elif opt.dataset == 'cifar100':
        nbclasses=100
        print('Dataset : Cifar100')
    # bsz: batch_size
    #output = torch.randn(batch_size, features)
   
    output=torch.randn(features.shape)
                         
    targets = torch.randint(0, nbclasses, (bsz,)).cuda(non_blocking=True)
        
    class_features = {idx: [features] for idx in range(nbclasses)}
    
       
    for class_index in range(nbclasses):
        idx = targets == class_index
        class_feat = output[idx]
        class_features[class_index].extend(class_feat)
        
    loss = criterion(class_features, labels)