Update app.py
Browse files
app.py
CHANGED
@@ -21,7 +21,7 @@ class Generator(nn.Module):
|
|
21 |
return self.model(x)
|
22 |
device = "cpu"
|
23 |
generator = Generator()
|
24 |
-
generator.load_state_dict(torch.load('generator_model.pt'))
|
25 |
num_faces_to_generate = 10
|
26 |
z_dim = 13
|
27 |
ethnicity_map = {'White': 0, 'Black': 1, 'Asian': 2, 'Indian': 3, 'Other': 4}
|
|
|
21 |
return self.model(x)
|
22 |
device = "cpu"
|
23 |
generator = Generator()
|
24 |
+
generator.load_state_dict(torch.load('generator_model.pt', map_location=torch.device('cpu')))
|
25 |
num_faces_to_generate = 10
|
26 |
z_dim = 13
|
27 |
ethnicity_map = {'White': 0, 'Black': 1, 'Asian': 2, 'Indian': 3, 'Other': 4}
|