@Navpreet_Singh
You can define a mask (a matrix of only 0 & 1). Multiply the mask with the matrix you have & then take the sum. Example:
input_matrix  = torch.ones((4, 5, 5)) #batch_size is 4
mask = torch.zeros((4, 5, 5))
mask[:, torch.arange(0,5), torch.arange(0,5) ] = 1.0 #This will mask all non-diagonal values.

output = input_matrix * mask
output = torch.sum(output,axis=(1,2))  #output will be of dimension (batch) & will contain the trace of each 5 x 5 matrix.