@wuhoo The softmax activation needs in your case more than one output channel to calculate the probabilities properly.
See this short code snippet:
a = torch.FloatTensor(1, 1, 2, 2).normal_()
print a
>> (0 ,0 ,.,.) = 
>> 0.7764 -0.8651
>> 0.4080  0.9525
>> [torch.FloatTensor of size 1x1x2x2]

F.softmax(a) 
>> Variable containing:
>> (0 ,0 ,.,.) = 
>>  1  1
>> 1  1
>> [torch.FloatTensor of size 1x1x2x2]

F.sigmoid(a)
>> Variable containing:
>> (0 ,0 ,.,.) = 
>> 0.6849  0.2963
>> 0.6006  0.7216
>> [torch.FloatTensor of size 1x1x2x2]

As you can see, the softmax calculates the probabilities in dim 1.
Try this example with 2 channels in dim 1 and it should work as expected:
a = torch.FloatTensor(1, 2, 2, 2).normal_()
F.softmax(a)
>> Variable containing:
>> (0 ,0 ,.,.) = 
>> 0.5941  0.2936
>> 0.3891  0.3968

>> (0 ,1 ,.,.) = 
>> 0.4059  0.7064
>> 0.6109  0.6032
>> [torch.FloatTensor of size 1x2x2x2]