apunkt/ai-privacy-toolkit
1
0
Fork 0
mirror of https://github.com/IBM/ai-privacy-toolkit.git synced 2026-07-02 16:01:00 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge with main

Browse source
This commit is contained in:
abigailt 2022-08-01 18:12:03 +03:00
parent dc5cc793ee
commit 64038f76f9
2 changed files with 5 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

4
tests/test_pytorch.py
View file

@ -54,7 +54,7 @@ def test_nursery_pytorch_state_dict():
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(inner_model.parameters(), lr=0.01)
model = PyTorchClassifier(model=inner_model, output_type=ModelOutputType.CLASSIFIER_VECTOR, loss=criterion,
model = PyTorchClassifier(model=inner_model, output_type=ModelOutputType.CLASSIFIER_LOGITS, loss=criterion,
optimizer=optimizer, input_shape=(24,),
nb_classes=4)
model.fit(PytorchData(x_train.astype(np.float32), y_train), save_entire_model=False, nb_epochs=10)
@ -84,7 +84,7 @@ def test_nursery_pytorch_save_entire_model():
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
art_model = PyTorchClassifier(model=model, output_type=ModelOutputType.CLASSIFIER_VECTOR, loss=criterion,
art_model = PyTorchClassifier(model=model, output_type=ModelOutputType.CLASSIFIER_LOGITS, loss=criterion,
optimizer=optimizer, input_shape=(24,),
nb_classes=4)
art_model.fit(PytorchData(x_train.astype(np.float32), y_train), save_entire_model=True, nb_epochs=10)