Ok, I initialized all my Linear weights based on this comment 6, but pytorch is still not learning:
def weight_init(m): 
  if isinstance(m, nn.Linear):
    size = m.weight.size()
    fan_out = size[0] # number of rows
    fan_in = size[1] # number of columns
    variance = np.sqrt(2.0/(fan_in + fan_out))
    m.weight.data.normal_(0.0, variance)
    
baseline_model = FCBaseline(SEGMENT_WIDTH, SEGMENT_HEIGHT, SEGMENT_CHANNELS, num_classes)
baseline_model.apply(weight_init)