Update README.md

Fix the problem of missing definition of device

README.md

@@ -108,6 +108,7 @@ class EnsembleModel(nn.Module):
         output = self.fc(torch.cat((convnext_out, mobilenet_out, efficientnet_out), dim=1))
         return output
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 convnext_model = CustomConvNeXtModel(weights=ConvNeXt_Tiny_Weights.DEFAULT, num_classes=2)
 mobilenet_model = CustomMobileNetModel(weights=MobileNet_V2_Weights.DEFAULT, num_classes=2)
 efficientnet_model = CustomEfficientNetModel(weights=EfficientNet_B0_Weights.DEFAULT, num_classes=2)