Fix code style in PyTorch QuickStart (#4821)
* fix two issues * fix two issues
This commit is contained in:
parent
e7c7ab2b8e
commit
04542bc6db
1 changed files with 3 additions and 2 deletions
|
|
@ -26,6 +26,7 @@ pip install jep==3.9.0
|
||||||
```python
|
```python
|
||||||
from bigdl.orca import init_orca_context, stop_orca_context
|
from bigdl.orca import init_orca_context, stop_orca_context
|
||||||
|
|
||||||
|
cluster_mode = "local"
|
||||||
if cluster_mode == "local": # For local machine
|
if cluster_mode == "local": # For local machine
|
||||||
init_orca_context(cores=4, memory="10g")
|
init_orca_context(cores=4, memory="10g")
|
||||||
elif cluster_mode == "k8s": # For K8s cluster
|
elif cluster_mode == "k8s": # For K8s cluster
|
||||||
|
|
@ -59,7 +60,7 @@ class LeNet(nn.Module):
|
||||||
self.conv2 = nn.Conv2d(20, 50, 5, 1)
|
self.conv2 = nn.Conv2d(20, 50, 5, 1)
|
||||||
self.fc1 = nn.Linear(4*4*50, 500)
|
self.fc1 = nn.Linear(4*4*50, 500)
|
||||||
self.fc2 = nn.Linear(500, 10)
|
self.fc2 = nn.Linear(500, 10)
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
x = F.relu(self.conv1(x))
|
x = F.relu(self.conv1(x))
|
||||||
x = F.max_pool2d(x, 2, 2)
|
x = F.max_pool2d(x, 2, 2)
|
||||||
|
|
@ -69,7 +70,7 @@ class LeNet(nn.Module):
|
||||||
x = F.relu(self.fc1(x))
|
x = F.relu(self.fc1(x))
|
||||||
x = self.fc2(x)
|
x = self.fc2(x)
|
||||||
return F.log_softmax(x, dim=1)
|
return F.log_softmax(x, dim=1)
|
||||||
|
|
||||||
model = LeNet()
|
model = LeNet()
|
||||||
model.train()
|
model.train()
|
||||||
criterion = nn.NLLLoss()
|
criterion = nn.NLLLoss()
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue