human
I am working on the file "Sessions/Session14/Day2/DeeplearningSolutions.ipynb".
The first task is:
Challenge Problem Add a third class to your classifier and begin accounting for uneven classes. There are several steps to this:

Edit the neural network to output 3 classes
Change the criterion to a custom criterion function, such that the entropy of each class is weighted by the inverse fraction of each class size (e.g., if the galaxy class breakdowns are 1:2:3, the weights would be 6:3:2).
Can you write Python code for it?
 ... Give me a the Python code you requested?
assistant
#Here is the Python code you requested.

class MLP_new(torch.nn.Module):
      # this defines the model
        def __init__(self, input_size, hidden_size):
            super(MLP_new, self).__init__()
            print(input_size,hidden_size)
            self.input_size = input_size
            self.hidden_size  = hidden_size
            self.hiddenlayer = torch.nn.Linear(self.input_size, self.hidden_size)
            self.outputlayer = torch.nn.Linear(self.hidden_size, 3)
            self.sigmoid = torch.nn.Sigmoid()
            self.softmax = torch.nn.Softmax()
        def forward(self, x):
            layer1 = self.hiddenlayer(x)
            activation = self.sigmoid(layer1)
            layer2 = self.outputlayer(activation)
            activation2 = self.sigmoid(layer1)
            layer3 = self.outputlayer(activation2)
            output = self.softmax(layer3)
            return output

#Only work with 0,1,2
gind = np.where((labels_original==0) | (labels_original==1) | (labels_original==2))
images_top_three = images[gind]
labels_top_three = labels_original[gind]

x,counts = np.unique(labels_top_three,return_counts=True)
print(counts)

torch.set_default_dtype(torch.float)
labels_top_three_one_hot = F.one_hot(torch.tensor(labels_top_three - np.min(labels_top_three)).long(), num_classes=3)
images_top_three = torch.tensor(images_top_three).float()
labels_top_three_one_hot = labels_top_three_one_hot.float()
# we're going to flatten the images for our MLP
images_top_three_flat = images_top_three.reshape(len(images_top_three),-1)

#Normalize the flux of the images here
images_top_three_flat = (images_top_three_flat  - torch.mean(images_top_three_flat))/torch.std(images_top_three_flat)
images_train, images_test, labels_train, labels_test = train_test_split(
    images_top_three_flat, labels_top_three_one_hot, test_size=0.33, random_state=42)

# train the model
def train_model(training_data,training_labels, test_data,test_labels, model):
  # define the optimization
  criterion = torch.nn.CrossEntropyLoss(weight=torch.Tensor(np.sum(counts)/counts))
 ...