Here is the psudue code of what I do:
Let’s say I am loading only three layers of the network and each has two sublayers called l1 and l2:
for i in range (3):
self.layer[i]…l1.weight.copy_(state_dict[“fc_{}.l1.weight”.format(i)], non_blocking=True)
self.layer[i].l2.weight.copy_(state_dict[“fc_{}.l2.weight”.format(i)], non_blocking=True)
This should not be related to batch size, but with different batch sizes, I get different timing. I use time.time from python time package to calculate the timing. Also, I am copying from cpu to gpu (my state_dict is on cpu and my model is on gpu).