Spaces:
Sleeping
Sleeping
Update model.py
Browse files
model.py
CHANGED
@@ -1,7 +1,7 @@
|
|
1 |
"""
|
2 |
Implementation of YOLOv3 architecture
|
3 |
"""
|
4 |
-
|
5 |
import torch
|
6 |
import torch.nn as nn
|
7 |
|
@@ -100,7 +100,7 @@ class ScalePrediction(nn.Module):
|
|
100 |
)
|
101 |
|
102 |
|
103 |
-
class YOLOv3(
|
104 |
def __init__(self, in_channels=3, num_classes=80):
|
105 |
super().__init__()
|
106 |
self.num_classes = num_classes
|
|
|
1 |
"""
|
2 |
Implementation of YOLOv3 architecture
|
3 |
"""
|
4 |
+
from pytorch_lightning import LightningModule
|
5 |
import torch
|
6 |
import torch.nn as nn
|
7 |
|
|
|
100 |
)
|
101 |
|
102 |
|
103 |
+
class YOLOv3(LightningModule):
|
104 |
def __init__(self, in_channels=3, num_classes=80):
|
105 |
super().__init__()
|
106 |
self.num_classes = num_classes
|