I am trying to use resnet50 quantized model for transfer learning and training the last fully connected layer. From the documentation what I figured was quantized model can only be trained on CPUs but the non-quantized layer can be trained on GPU but how to transfer this layer to GPU? Cannot find much related to this anywhere.
Model looks like:
def create_combined_model(model_fe): # Step 1. Isolate the feature extractor. model_fe_features = nn.Sequential( model_fe.quant, # Quantize the input model_fe.conv1, model_fe.bn1, model_fe.relu, model_fe.maxpool, model_fe.layer1, model_fe.layer2, model_fe.layer3, model_fe.layer4, model_fe.avgpool, model_fe.dequant, # Dequantize the output ) # Step 2. Create a new "head" new_head = nn.Sequential( nn.Dropout(p=0.5), nn.Linear(num_ftrs, num_classes), ) new_model = nn.Sequential( model_fe_features, nn.Flatten(1), new_head, ) return new_model