So this worked for me which is little weird of a workflow:
# While saving checkpoint i.e. comment out while loading checkpoint
    if torch.cuda.device_count() > 1:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
        network = nn.DataParallel(network)
    
    
    network.to(device)
    optimizer = optim.AdamW(network.parameters(), lr=run.lr, weight_decay=run.weight_decay)
    
#     try:
    with open('check-point.pth', 'rb') as f:
        print('file opened')
        checkpoint = torch.load(f)
        print('file loaded')
        network.load_state_dict(checkpoint['model_state_dict'])
        print('network loaded')
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        print('optimizer loaded')
        epoc = checkpoint['epoch']
        print(f'blah epoc: {epoc}')    
    # While loading checkpoint
        if torch.cuda.device_count() > 1:
            print("Let's use", torch.cuda.device_count(), "GPUs!")
            network = nn.DataParallel(network)